From 81973c73b24c1103bc3e1145fb6aa11063f17086 Mon Sep 17 00:00:00 2001 From: Som Tripathi Date: Tue, 21 Jul 2026 23:28:56 -0500 Subject: [PATCH 01/15] Add message-specific regression test for AdagradDA grad-rank validation Addresses #94130. Master already validates that grad's rank matches var's rank in the Sparse*Op family (landed separately, covers ResourceSparseApplyAdagradDA among others), but the existing regression test (testSparseApplyOpsRejectLowerRankGrad) only asserts on the exception type, not the message. Per gemini-code-assist review feedback, assertRaises alone could pass even if a different InvalidArgumentError fired for an unrelated reason. Added testResourceSparseApplyAdagradDAInvalidGradRank, which pins the specific "grad must have the same number of dimensions as var" message via assertRaisesRegex for this op, so the test can only pass if the intended validation actually fired. --- .../python/training/training_ops_test.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/tensorflow/python/training/training_ops_test.py b/tensorflow/python/training/training_ops_test.py index b8a6106c2a1e86..eaa72784d43c18 100644 --- a/tensorflow/python/training/training_ops_test.py +++ b/tensorflow/python/training/training_ops_test.py @@ -576,6 +576,35 @@ def testSparseApplyOpsRejectLowerRankGrad(self): with self.assertRaises(errors.InvalidArgumentError): self.evaluate(apply_op()) + @test_util.run_v2_only + def testResourceSparseApplyAdagradDAInvalidGradRank(self): + # A scalar `grad` with a higher-rank `var` used to hit a fatal CHECK + # instead of raising InvalidArgumentError (see GitHub issue #94130). + var = variables.Variable([[0.0, 0.0]] * 10, dtype=dtypes.float32) + gradient_accumulator = variables.Variable( + [[0.0, 0.0]] * 10, dtype=dtypes.float32) + gradient_squared_accumulator = variables.Variable( + [[0.0, 0.0]] * 10, dtype=dtypes.float32) + self.evaluate(variables.global_variables_initializer()) + + grad = constant_op.constant(0.0, dtype=dtypes.float32) # wrong rank + indices = constant_op.constant([0, 0], dtype=dtypes.int32) + + with self.assertRaisesRegex( + errors.InvalidArgumentError, + "grad must have the same number of dimensions as var"): + self.evaluate( + gen_training_ops.resource_sparse_apply_adagrad_da( + var.handle, + gradient_accumulator.handle, + gradient_squared_accumulator.handle, + grad, + indices, + constant_op.constant(0.0, dtype=dtypes.float32), + constant_op.constant(0.0, dtype=dtypes.float32), + constant_op.constant(0.0, dtype=dtypes.float32), + constant_op.constant(1, dtype=dtypes.int64))) + if __name__ == '__main__': googletest.main() From c0d9f964295f3667ff945dd748cae4b97960b382 Mon Sep 17 00:00:00 2001 From: Som Tripathi Date: Thu, 23 Jul 2026 03:00:15 -0500 Subject: [PATCH 02/15] Pin ResourceSparseApplyAdagradDA test to CPU: no GPU kernel exists CI on real GPU hardware failed: ResourceSparseApplyAdagradDA has no registered GPU kernel (CPU only), so without explicit device placement the op's variables default onto GPU:0 and the test fails with a device placement error before ever reaching the validation code under test. Wrapped the test body in `with ops.device("/cpu:0")` so it exercises the CPU kernel deterministically regardless of what hardware CI runs on. --- .../python/training/training_ops_test.py | 49 ++++++++++--------- 1 file changed, 26 insertions(+), 23 deletions(-) diff --git a/tensorflow/python/training/training_ops_test.py b/tensorflow/python/training/training_ops_test.py index eaa72784d43c18..38faca094f0fde 100644 --- a/tensorflow/python/training/training_ops_test.py +++ b/tensorflow/python/training/training_ops_test.py @@ -580,30 +580,33 @@ def testSparseApplyOpsRejectLowerRankGrad(self): def testResourceSparseApplyAdagradDAInvalidGradRank(self): # A scalar `grad` with a higher-rank `var` used to hit a fatal CHECK # instead of raising InvalidArgumentError (see GitHub issue #94130). - var = variables.Variable([[0.0, 0.0]] * 10, dtype=dtypes.float32) - gradient_accumulator = variables.Variable( - [[0.0, 0.0]] * 10, dtype=dtypes.float32) - gradient_squared_accumulator = variables.Variable( - [[0.0, 0.0]] * 10, dtype=dtypes.float32) - self.evaluate(variables.global_variables_initializer()) + # ResourceSparseApplyAdagradDA only has a CPU kernel, so pin the whole + # test to CPU rather than relying on default placement. + with ops.device("/cpu:0"): + var = variables.Variable([[0.0, 0.0]] * 10, dtype=dtypes.float32) + gradient_accumulator = variables.Variable( + [[0.0, 0.0]] * 10, dtype=dtypes.float32) + gradient_squared_accumulator = variables.Variable( + [[0.0, 0.0]] * 10, dtype=dtypes.float32) + self.evaluate(variables.global_variables_initializer()) - grad = constant_op.constant(0.0, dtype=dtypes.float32) # wrong rank - indices = constant_op.constant([0, 0], dtype=dtypes.int32) - - with self.assertRaisesRegex( - errors.InvalidArgumentError, - "grad must have the same number of dimensions as var"): - self.evaluate( - gen_training_ops.resource_sparse_apply_adagrad_da( - var.handle, - gradient_accumulator.handle, - gradient_squared_accumulator.handle, - grad, - indices, - constant_op.constant(0.0, dtype=dtypes.float32), - constant_op.constant(0.0, dtype=dtypes.float32), - constant_op.constant(0.0, dtype=dtypes.float32), - constant_op.constant(1, dtype=dtypes.int64))) + grad = constant_op.constant(0.0, dtype=dtypes.float32) # wrong rank + indices = constant_op.constant([0, 0], dtype=dtypes.int32) + + with self.assertRaisesRegex( + errors.InvalidArgumentError, + "grad must have the same number of dimensions as var"): + self.evaluate( + gen_training_ops.resource_sparse_apply_adagrad_da( + var.handle, + gradient_accumulator.handle, + gradient_squared_accumulator.handle, + grad, + indices, + constant_op.constant(0.0, dtype=dtypes.float32), + constant_op.constant(0.0, dtype=dtypes.float32), + constant_op.constant(0.0, dtype=dtypes.float32), + constant_op.constant(1, dtype=dtypes.int64))) if __name__ == '__main__': From 70d3e223c910d3abdf2186dc6427bd8255cd8747 Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Thu, 13 Aug 2026 16:20:04 -0400 Subject: [PATCH 03/15] [XLA] Fix RandomShuffle comparison type for float vectors --- .../compiler/mlir/tf2xla/tests/legalize-tf.mlir | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir b/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir index be0bb7167f2863..195c1fadfad69b 100644 --- a/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir +++ b/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir @@ -2708,6 +2708,17 @@ func.func @neg_dynamic(%arg0: tensor) -> tensor { // ----- +// CHECK-LABEL: @random_shuffle_float_vector +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]] : (tensor, tensor) -> tensor + %0 = "tf.RandomShuffle"(%arg0) {seed = 1 : i64, seed2 = 2 : i64} : (tensor<2xf32>) -> tensor<2xf32> + func.return %0 : tensor<2xf32> +} + +// ----- + // CHECK-LABEL: @sigmoid func.func @sigmoid(%arg0: tensor<2xf32>) -> tensor<2xf32> { // CHECK: mhlo.logistic @@ -2798,4 +2809,3 @@ func.func @func_xla_sharding_consistent(%arg0: tensor<4x8xi32>) -> (tensor<4x8xi %1 = "tf.A"(%0) : (tensor<4x8xi32>) -> (tensor<4x8xi32>) func.return %1 : tensor<4x8xi32> } - From 6eac4b01b9cd1b5dc527f9d6814fbb8687d6ac13 Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Tue, 1 Sep 2026 05:08:52 -0400 Subject: [PATCH 04/15] Update RandomShuffle comparator check Match the TOTALORDER comparison attribute emitted by the current tf2xla legalization. --- tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir b/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir index 195c1fadfad69b..a4d4f6a6ced441 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]] : (tensor, tensor) -> tensor + // CHECK: mhlo.compare LT, [[KEY_LHS]], [[KEY_RHS]], TOTALORDER : (tensor, tensor) -> tensor %0 = "tf.RandomShuffle"(%arg0) {seed = 1 : i64, seed2 = 2 : i64} : (tensor<2xf32>) -> tensor<2xf32> func.return %0 : tensor<2xf32> } From 69e807c19df417aa8ddaf4d5eefc38c3306305bf Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Wed, 12 Aug 2026 13:56:00 -0400 Subject: [PATCH 05/15] Fix SELU gradient precision for negative inputs --- .../eager/pywrap_gradient_exclusions.cc | 6 +++--- tensorflow/python/ops/nn_grad.py | 15 +++++++++++++- tensorflow/python/ops/nn_grad_test.py | 20 +++++++++++++++++++ 3 files changed, 37 insertions(+), 4 deletions(-) diff --git a/tensorflow/python/eager/pywrap_gradient_exclusions.cc b/tensorflow/python/eager/pywrap_gradient_exclusions.cc index 6e39edcad17aa1..3ff8a8ce366a1c 100644 --- a/tensorflow/python/eager/pywrap_gradient_exclusions.cc +++ b/tensorflow/python/eager/pywrap_gradient_exclusions.cc @@ -50,7 +50,7 @@ auto OpGradientInfoInit(const T &a) { absl::optional> OpGradientUnusedInputIndices( const tensorflow::string &op_name) { - static std::array a = {{ + static std::array a = {{ {"Acosh"}, {"AllToAll", 1, {0}}, {"ApproximateEqual"}, @@ -292,7 +292,6 @@ absl::optional> OpGradientUnusedInputIndices( {"SdcaFprint"}, {"SegmentSum", 1, {0}}, {"Select", 1, {2}}, - {"Selu"}, {"SerializeTensor"}, {"SetSize"}, {"Shape"}, @@ -429,7 +428,7 @@ absl::optional> OpGradientUnusedInputIndices( absl::optional> OpGradientUnusedOutputIndices( const tensorflow::string &op_name) { - static std::array a = {{ + static std::array a = {{ {"Abs"}, {"AccumulateNV2"}, {"Acos"}, @@ -760,6 +759,7 @@ absl::optional> OpGradientUnusedOutputIndices( {"SegmentMean"}, {"SegmentSum"}, {"Select"}, + {"Selu"}, {"SeluGrad"}, {"SerializeTensor"}, {"SetSize"}, diff --git a/tensorflow/python/ops/nn_grad.py b/tensorflow/python/ops/nn_grad.py index 1b7dcc0ac9314a..8690004f114cd3 100644 --- a/tensorflow/python/ops/nn_grad.py +++ b/tensorflow/python/ops/nn_grad.py @@ -18,6 +18,7 @@ import itertools import operator +from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops @@ -466,7 +467,19 @@ def _EluGrad(op: ops.Operation, grad): @ops.RegisterGradient("Selu") def _SeluGrad(op: ops.Operation, grad): - return gen_nn_ops.selu_grad(grad, op.outputs[0]) + x = op.inputs[0] + scale = constant_op.constant( + 1.0507009873554804934193349852946, dtype=x.dtype) + scale_alpha = constant_op.constant( + 1.7580993408473768599402175208123, dtype=x.dtype) + # Reconstructing the negative-branch derivative from the SELU output loses + # precision when the output rounds to -scale_alpha. Compute it from x. + derivative = array_ops.where_v2( + x < 0.0, + scale_alpha * math_ops.exp(math_ops.minimum(x, 0.0)), + scale) + derivative = array_ops.where_v2(math_ops.is_nan(x), x, derivative) + return grad * derivative @ops.RegisterGradient("Softplus") diff --git a/tensorflow/python/ops/nn_grad_test.py b/tensorflow/python/ops/nn_grad_test.py index 1f2b82ecbe9dea..b59c9fe47b7b21 100644 --- a/tensorflow/python/ops/nn_grad_test.py +++ b/tensorflow/python/ops/nn_grad_test.py @@ -237,6 +237,26 @@ def testEluGradGradWRTinputs(self): class SeluGradOpTest(test.TestCase): + @test_util.run_in_graph_and_eager_modes + def testSeluGradPreservesSmallNegativeGradients(self): + scale_alpha = 1.7580993408473768599402175208123 + test_cases = ( + (dtypes.float32, -20.0, 1e-6), + (dtypes.float64, -700.0, 1e-14), + ) + for dtype, value, rtol in test_cases: + with self.subTest(dtype=dtype.name): + inputs = constant_op.constant(value, dtype=dtype) + with backprop.GradientTape() as tape: + tape.watch(inputs) + selu = gen_nn_ops.selu(inputs) + + selu_grad = tape.gradient(selu, inputs) + np_dtype = dtype.as_numpy_dtype + expected = np_dtype(scale_alpha) * np.exp(np_dtype(value)) + self.assertAllClose( + expected, self.evaluate(selu_grad), rtol=rtol, atol=0) + @test_util.run_deprecated_v1 def testSeluGradGradWRTgrad_ys(self): inputs = constant_op.constant( From 2986a870f07e5380e45de6f52b4429689b1e9c85 Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Wed, 12 Aug 2026 14:11:22 -0400 Subject: [PATCH 06/15] Add SELU gradient edge-case coverage --- tensorflow/python/ops/nn_grad_test.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tensorflow/python/ops/nn_grad_test.py b/tensorflow/python/ops/nn_grad_test.py index b59c9fe47b7b21..976d82e5dae4b3 100644 --- a/tensorflow/python/ops/nn_grad_test.py +++ b/tensorflow/python/ops/nn_grad_test.py @@ -257,6 +257,33 @@ def testSeluGradPreservesSmallNegativeGradients(self): self.assertAllClose( expected, self.evaluate(selu_grad), rtol=rtol, atol=0) + @test_util.run_in_graph_and_eager_modes + def testSeluGradEdgeCases(self): + scale = 1.0507009873554804934193349852946 + scale_alpha = 1.7580993408473768599402175208123 + for dtype in (dtypes.float32, dtypes.float64): + with self.subTest(dtype=dtype.name): + np_dtype = dtype.as_numpy_dtype + values = np.array( + [-np.inf, -2.0, -0.0, 0.0, 2.0, np.inf, np.nan], + dtype=np_dtype) + inputs = constant_op.constant(values, dtype=dtype) + with backprop.GradientTape() as tape: + tape.watch(inputs) + selu = gen_nn_ops.selu(inputs) + + selu_grad = self.evaluate(tape.gradient(selu, inputs)) + expected = np.array([ + 0.0, + np_dtype(scale_alpha) * np.exp(np_dtype(-2.0)), + scale, + scale, + scale, + scale, + np.nan, + ], dtype=np_dtype) + self.assertAllClose(expected, selu_grad) + @test_util.run_deprecated_v1 def testSeluGradGradWRTgrad_ys(self): inputs = constant_op.constant( From aa87766b576ccf66e2dee3545ed5ef04aaf2f41b Mon Sep 17 00:00:00 2001 From: Kuldeeep18 Date: Tue, 1 Sep 2026 15:37:32 +0530 Subject: [PATCH 07/15] Fix tnp.eye column count when M=0 --- tensorflow/python/ops/numpy_ops/np_array_ops.py | 2 +- tensorflow/python/ops/numpy_ops/np_array_ops_test.py | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py index 85d12f133b612e..5b7d87b83138af 100644 --- a/tensorflow/python/ops/numpy_ops/np_array_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py @@ -100,7 +100,7 @@ def ones_like(a, dtype=None): def eye(N, M=None, k=0, dtype=float): # pylint: disable=invalid-name,missing-docstring if dtype: dtype = np_utils.result_type(dtype) - if not M: + if M is None: M = N # Making sure N, M and k are `int` N = int(N) diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py index 1045b154b4d36e..21b055b2e701e7 100644 --- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py +++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py @@ -230,6 +230,13 @@ def testEye(self): np_array_ops.eye(n, m, k, dtype=dtype), np.eye(n, m, k, dtype=dtype)) + # Test M=0 and N=0 zero-dimension edge cases + for n in (0, 1, 3): + for m in (0, 1, 3): + self.match(np_array_ops.eye(n, m), np.eye(n, m)) + for k in range(-n - 1, m + 2): + self.match(np_array_ops.eye(n, m, k), np.eye(n, m, k)) + def testIdentity(self): n_max = 3 From 04ce8e9cba44ffa83a32385b23593ad47cac59f6 Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Tue, 1 Sep 2026 14:07:52 -0400 Subject: [PATCH 08/15] Add constant_op dependency to nn_grad --- tensorflow/python/ops/BUILD | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD index fea13cbd711535..126b26064c1cfd 100644 --- a/tensorflow/python/ops/BUILD +++ b/tensorflow/python/ops/BUILD @@ -2208,6 +2208,7 @@ py_library( ":array_ops_stack", ":math_ops", ":nn_ops_gen", + "//tensorflow/python/framework:constant_op", "//tensorflow/python/framework:dtypes", "//tensorflow/python/framework:ops", ], From 6d821817a5c026a281be9b1642377a24af11299e Mon Sep 17 00:00:00 2001 From: Hyeontaek Lim Date: Wed, 2 Sep 2026 17:34:14 -0700 Subject: [PATCH 09/15] [IFRT] Add `Sharding{Spec}::UniqueIndexDomains()` and `ShardToUniqueIndexDomainIndex()` For host-buffer transfers (such as `xla::ifrt::Client::CopyArraysToHostBufferShards`), the caller must identify the unique `IndexDomain`s for the array sharding and map them to corresponding shard indices. Currently, the user only has access to `xla::ifrt::Sharding{Spec}::IndexDomains()`, which returns one index domain per device shard without deduplication or native shard-to-unique-domain grouping. This is often expensive (`O(#devices)`) and is not in the form of the information that the user actually needs to use for the downstream API call. This change adds new sharding APIs: - `xla::ifrt::Sharding{Spec}::UniqueIndexDomains()`: Returns a vector of `IndexDomainAndShardIndices`, where each item contains a unique `IndexDomain` and an `absl::Span` view of shard indices mapped to that domain. - `xla::ifrt::Sharding{Spec}::ShardToUniqueIndexDomainIndex()`: Returns an `absl::Span` mapping each global shard index to the index of its corresponding unique index domain. The shard indices view and reverse mapping are cached independently inside the `ShardingSpec` instance and computed on-demand. The unique index domains are always ordered deterministically in a manner agnostic of shapes and device lists, but specific to the sharding type for optimal traversal and consistent use across APIs. PiperOrigin-RevId: 975419584 --- third_party/xla/xla/python/ifrt/BUILD | 1 + third_party/xla/xla/python/ifrt/sharding.cc | 11 + third_party/xla/xla/python/ifrt/sharding.h | 25 +- .../xla/xla/python/ifrt/sharding_spec.cc | 237 ++++++++++++++++++ .../xla/xla/python/ifrt/sharding_spec.h | 91 +++++++ .../xla/xla/python/ifrt/sharding_spec_test.cc | 140 +++++++++++ .../xla/xla/python/ifrt/sharding_test.cc | 109 ++++++++ third_party/xla/xla/python/pjrt_ifrt/BUILD | 2 + .../xla/xla/python/pjrt_ifrt/xla_sharding.cc | 1 + .../xla/xla/python/pjrt_ifrt/xla_sharding.h | 3 - .../xla/python/pjrt_ifrt/xla_sharding_spec.cc | 143 ++++++++++- .../xla/python/pjrt_ifrt/xla_sharding_spec.h | 14 ++ .../pjrt_ifrt/xla_sharding_spec_test.cc | 99 ++++++++ .../xla/python/pjrt_ifrt/xla_sharding_test.cc | 70 +++++- 14 files changed, 935 insertions(+), 11 deletions(-) diff --git a/third_party/xla/xla/python/ifrt/BUILD b/third_party/xla/xla/python/ifrt/BUILD index 7a5c8592aa348b..23bc8a98a5d733 100644 --- a/third_party/xla/xla/python/ifrt/BUILD +++ b/third_party/xla/xla/python/ifrt/BUILD @@ -126,6 +126,7 @@ cc_library( "//xla/tsl/platform:logging", "//xla/tsl/platform:statusor", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/base", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/base:no_destructor", "@com_google_absl//absl/base:nullability", diff --git a/third_party/xla/xla/python/ifrt/sharding.cc b/third_party/xla/xla/python/ifrt/sharding.cc index dd16a79114abae..a641d51f29e315 100644 --- a/third_party/xla/xla/python/ifrt/sharding.cc +++ b/third_party/xla/xla/python/ifrt/sharding.cc @@ -31,6 +31,7 @@ limitations under the License. #include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" +#include "absl/status/statusor.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" @@ -75,6 +76,16 @@ bool Sharding::operator==(const Sharding& other) const { *devices() == *other.devices(); } +absl::StatusOr> +Sharding::UniqueIndexDomains(const Shape& shape) const { + return sharding_spec()->UniqueIndexDomains(shape); +} + +absl::StatusOr> Sharding::ShardToUniqueIndexDomainIndex() + const { + return sharding_spec()->ShardToUniqueIndexDomainIndex(); +} + absl::StatusOr Sharding::FromProto( Client* client, const ShardingProto& sharding_proto) { return Deserialize( diff --git a/third_party/xla/xla/python/ifrt/sharding.h b/third_party/xla/xla/python/ifrt/sharding.h index 0b658604161553..7490a43a8d4656 100644 --- a/third_party/xla/xla/python/ifrt/sharding.h +++ b/third_party/xla/xla/python/ifrt/sharding.h @@ -21,13 +21,17 @@ limitations under the License. #include #include #include -#include #include +#include "absl/base/attributes.h" #include "absl/base/nullability.h" +#include "absl/container/inlined_vector.h" #include "absl/hash/hash.h" #include "absl/log/check.h" +#include "absl/status/status.h" #include "absl/status/status_macros.h" +#include "absl/status/statusor.h" +#include "absl/types/span.h" #include "xla/python/ifrt/device.h" #include "xla/python/ifrt/device_list.h" #include "xla/python/ifrt/index_domain.h" @@ -149,6 +153,25 @@ class Sharding : public RTTIExtends { const Shape& shape, SingleDeviceShardSemantics single_device_shard_semantics) const = 0; + using IndexDomainAndShardIndices = ShardingSpec::IndexDomainAndShardIndices; + + // Breaks a shape up into unique `IndexDomain`s and the shard indices mapped + // to it. The result is calculated for all shards. + // + // The result is valid for the lifetime of this `Sharding`. + absl::StatusOr> + UniqueIndexDomains(const Shape& shape) const ABSL_ATTRIBUTE_LIFETIME_BOUND; + + // Inverse of `UniqueIndexDomains()` for `shard_indices`. Does not take + // `shape` because the result is independent of `shape`. + // + // Suppose `j` be `unique_index_domain_indices[shard_i]`. Then, + // `unique_index_domains[j].shard_indices` contains `shard_i`. + // + // The result is valid for the lifetime of this `Sharding`. + absl::StatusOr> ShardToUniqueIndexDomainIndex() const + ABSL_ATTRIBUTE_LIFETIME_BOUND; + template friend H AbslHashValue(H h, const Sharding& value) { value.Hash(absl::HashState::Create(&h)); diff --git a/third_party/xla/xla/python/ifrt/sharding_spec.cc b/third_party/xla/xla/python/ifrt/sharding_spec.cc index 72f3746521e526..a0d3b02d82f454 100644 --- a/third_party/xla/xla/python/ifrt/sharding_spec.cc +++ b/third_party/xla/xla/python/ifrt/sharding_spec.cc @@ -15,6 +15,7 @@ limitations under the License. #include "xla/python/ifrt/sharding_spec.h" +#include #include #include #include @@ -26,6 +27,8 @@ limitations under the License. #include #include "absl/algorithm/container.h" +#include "absl/base/call_once.h" +#include "absl/container/flat_hash_map.h" #include "absl/container/inlined_vector.h" #include "absl/hash/hash.h" #include "absl/log/check.h" @@ -244,6 +247,23 @@ absl::StatusOr> SingleDeviceShardingSpec::IndexDomains( return std::vector{IndexDomain(shape)}; } +absl::StatusOr> +SingleDeviceShardingSpec::UniqueIndexDomains(const Shape& shape) const { + static constexpr std::array kShardIndices({0}); + return absl::InlinedVector{ + IndexDomainAndShardIndices{ + /*index_domain=*/IndexDomain(shape), + /*shard_indices=*/absl::MakeConstSpan(kShardIndices), + }, + }; +} + +absl::StatusOr> +SingleDeviceShardingSpec::ShardToUniqueIndexDomainIndex() const { + static constexpr std::array kShardToUniqueIndexDomainIndex({0}); + return absl::MakeConstSpan(kShardToUniqueIndexDomainIndex); +} + std::string SingleDeviceShardingSpec::DebugString() const { return "SingleDeviceShardingSpec()"; } @@ -311,6 +331,18 @@ absl::StatusOr> OpaqueShardingSpec::IndexDomains( "OpaqueShardingSpec does not have index domain information"); } +absl::StatusOr> +OpaqueShardingSpec::UniqueIndexDomains(const Shape& shape) const { + return absl::InvalidArgumentError( + "OpaqueShardingSpec does not support UniqueIndexDomains"); +} + +absl::StatusOr> +OpaqueShardingSpec::ShardToUniqueIndexDomainIndex() const { + return absl::InvalidArgumentError( + "OpaqueShardingSpec does not support ShardToUniqueIndexDomainIndex"); +} + std::string OpaqueShardingSpec::DebugString() const { return absl::StrFormat("OpaqueShardingSpec(num_shards: %d)", num_shards_); } @@ -370,6 +402,13 @@ ConcreteShardingSpec::ConcreteShardingSpec( shape_(std::move(dynamic_shape)), shard_shapes_(std::move(shard_dynamic_shapes)) {} +ConcreteShardingSpec::ConcreteShardingSpec(const ConcreteShardingSpec& other) + : RTTIExtends(other), + shape_(other.shape_), + shard_shapes_(other.shard_shapes_), + shard_shape_(other.shard_shape_), + index_domains_(other.index_domains_) {} + absl::StatusOr ConcreteShardingSpec::ToSharding( DeviceListRef devices, MemoryKind memory_kind) const { if (devices->size() != num_shards()) { @@ -474,6 +513,76 @@ absl::StatusOr> ConcreteShardingSpec::IndexDomains( return *index_domains_; } +absl::StatusOr> +ConcreteShardingSpec::UniqueIndexDomains(const Shape& shape) const { + if (!index_domains_.has_value()) { + return absl::InvalidArgumentError( + "ConcreteShardingSpec does not have index domain information"); + } + if (has_static_shape() && this->shape() != shape) { + return absl::InvalidArgumentError(absl::StrFormat( + "ConcreteShardingSpec has index domains for shape %v, but was asked " + "to get unique index domains for shape %v", + this->shape(), shape)); + } + absl::call_once(unique_shard_indices_once_, [this] { + absl::flat_hash_map index_domain_to_unique_idx; + std::vector> shard_indices; + for (int i = 0; i < index_domains_->size(); ++i) { + const IndexDomain& domain = (*index_domains_)[i]; + auto [it, inserted] = + index_domain_to_unique_idx.try_emplace(domain, shard_indices.size()); + if (inserted) { + shard_indices.emplace_back(); + } + shard_indices[it->second].push_back(i); + } + cached_shard_indices_.reserve(index_domains_->size()); + cached_shard_indices_offsets_.reserve(shard_indices.size() + 1); + for (const auto& indices : shard_indices) { + cached_shard_indices_offsets_.push_back(cached_shard_indices_.size()); + cached_shard_indices_.insert(cached_shard_indices_.end(), indices.begin(), + indices.end()); + } + cached_shard_indices_offsets_.push_back(cached_shard_indices_.size()); + }); + + const int num_unique = cached_shard_indices_offsets_.size() - 1; + absl::InlinedVector unique_domains; + unique_domains.reserve(num_unique); + for (int i = 0; i < num_unique; ++i) { + const int offset = cached_shard_indices_offsets_[i]; + const int count = cached_shard_indices_offsets_[i + 1] - offset; + const int first_shard = cached_shard_indices_[offset]; + unique_domains.push_back(IndexDomainAndShardIndices{ + /*index_domain=*/(*index_domains_)[first_shard], + /*shard_indices=*/ + absl::MakeConstSpan(cached_shard_indices_).subspan(offset, count), + }); + } + return unique_domains; +} + +absl::StatusOr> +ConcreteShardingSpec::ShardToUniqueIndexDomainIndex() const { + if (!index_domains_.has_value()) { + return absl::InvalidArgumentError( + "ConcreteShardingSpec does not have index domain information"); + } + absl::call_once(shard_to_unique_index_domain_index_once_, [this] { + absl::flat_hash_map domain_to_unique_idx; + cached_shard_to_unique_index_domain_index_.reserve(index_domains_->size()); + for (int i = 0; i < index_domains_->size(); ++i) { + const IndexDomain& domain = (*index_domains_)[i]; + auto it = + domain_to_unique_idx.try_emplace(domain, domain_to_unique_idx.size()) + .first; + cached_shard_to_unique_index_domain_index_.push_back(it->second); + } + }); + return absl::MakeConstSpan(cached_shard_to_unique_index_domain_index_); +} + std::string ConcreteShardingSpec::DebugString() const { return std::visit( [this](const auto& shape, const auto& shard_shapes) { @@ -508,6 +617,12 @@ ConcreteEvenShardingSpec::ConcreteEvenShardingSpec(int num_shards, Shape shape, shape_(std::move(shape)), shard_shape_(std::move(shard_shape)) {} +ConcreteEvenShardingSpec::ConcreteEvenShardingSpec( + const ConcreteEvenShardingSpec& other) + : RTTIExtends(other), + shape_(other.shape_), + shard_shape_(other.shard_shape_) {} + absl::StatusOr ConcreteEvenShardingSpec::ToSharding( DeviceListRef devices, MemoryKind memory_kind) const { if (devices->size() != num_shards()) { @@ -584,6 +699,39 @@ absl::StatusOr> ConcreteEvenShardingSpec::IndexDomains( "ConcreteEvenShardingSpec does not have index domain information"); } +absl::StatusOr> +ConcreteEvenShardingSpec::UniqueIndexDomains(const Shape& shape) const { + if (!IsFullyReplicated() || this->shape() != shard_shape() || + this->shape() != shape) { + return absl::InvalidArgumentError( + "ConcreteEvenShardingSpec does not have index domain information"); + } + absl::call_once(unique_shard_indices_once_, [this] { + cached_shard_indices_.reserve(num_shards_); + for (int i = 0; i < num_shards_; ++i) { + cached_shard_indices_.push_back(i); + } + }); + return absl::InlinedVector{ + IndexDomainAndShardIndices{ + /*index_domain=*/IndexDomain(shape), + /*shard_indices=*/absl::MakeConstSpan(cached_shard_indices_), + }, + }; +} + +absl::StatusOr> +ConcreteEvenShardingSpec::ShardToUniqueIndexDomainIndex() const { + if (!IsFullyReplicated() || this->shape() != shard_shape()) { + return absl::InvalidArgumentError( + "ConcreteEvenShardingSpec does not have index domain information"); + } + absl::call_once(shard_to_unique_index_domain_index_once_, [this] { + cached_shard_to_unique_index_domain_index_.assign(num_shards_, 0); + }); + return absl::MakeConstSpan(cached_shard_to_unique_index_domain_index_); +} + std::string ConcreteEvenShardingSpec::DebugString() const { return absl::StrFormat( "ConcreteEvenShardingSpec(num_shards: %d, shape: %v, " @@ -610,6 +758,11 @@ ShardingParamShardingSpec::ShardingParamShardingSpec( num_shards, ComputeIsFullyReplicated(sharding_param)), sharding_param_(std::move(sharding_param)) {} +ShardingParamShardingSpec::ShardingParamShardingSpec( + const ShardingParamShardingSpec& other) + : RTTIExtends(other), + sharding_param_(other.sharding_param_) {} + absl::StatusOr ShardingParamShardingSpec::ToSharding( DeviceListRef devices, MemoryKind memory_kind) const { if (devices->size() != num_shards()) { @@ -722,6 +875,90 @@ ShardingParamShardingSpec::IndexDomains(const Shape& shape) const { return result; } +absl::StatusOr> +ShardingParamShardingSpec::UniqueIndexDomains(const Shape& shape) const { + ABSL_ASSIGN_OR_RETURN(Shape local_shape, GetShardShape(shape)); + + absl::call_once(unique_shard_indices_once_, [this] { + absl::InlinedVector device_list; + sharding_param_.minor_to_major().ToDeviceList(device_list); + if (device_list.size() != num_shards_) { + cached_shard_indices_ = absl::InvalidArgumentError(absl::StrFormat( + "ShardingParamShardingSpec has %d shards, but sharding param has %d " + "shards", + num_shards_, device_list.size())); + return; + } + cached_shard_indices_ = + std::vector(device_list.begin(), device_list.end()); + }); + ABSL_RETURN_IF_ERROR(cached_shard_indices_.status()); + + std::vector tile_indices = + GetTileIndices(sharding_param_.dim_shards()); + const int num_unique_tiles = tile_indices.size(); + if (num_shards_ % num_unique_tiles != 0) { + return absl::InvalidArgumentError(absl::StrFormat( + "ShardingParamShardingSpec has %d shards, but sharding param has %d " + "unique tiles, which is not a divisor of the number of shards", + num_shards_, num_unique_tiles)); + } + const int replication = num_shards_ / num_unique_tiles; + + absl::InlinedVector unique_domains; + unique_domains.reserve(num_unique_tiles); + for (int tile_idx = 0; tile_idx < num_unique_tiles; ++tile_idx) { + const Index& tile_index = tile_indices[tile_idx]; + unique_domains.push_back(IndexDomainAndShardIndices{ + /*index_domain=*/IndexDomain(tile_index * local_shape.dims(), + local_shape), + /*shard_indices=*/ + absl::MakeConstSpan(*cached_shard_indices_) + .subspan(tile_idx * replication, replication), + }); + } + + return unique_domains; +} + +absl::StatusOr> +ShardingParamShardingSpec::ShardToUniqueIndexDomainIndex() const { + absl::call_once(shard_to_unique_index_domain_index_once_, [this] { + std::vector tile_indices = + GetTileIndices(sharding_param_.dim_shards()); + const int num_unique_tiles = tile_indices.size(); + absl::InlinedVector device_list; + sharding_param_.minor_to_major().ToDeviceList(device_list); + if (device_list.size() != num_shards_) { + cached_shard_to_unique_index_domain_index_ = absl::InvalidArgumentError( + absl::StrFormat("ShardingParamShardingSpec has %d shards, but " + "sharding param has %d shards", + num_shards_, device_list.size())); + return; + } + if (device_list.size() % num_unique_tiles != 0) { + cached_shard_to_unique_index_domain_index_ = + absl::InvalidArgumentError(absl::StrFormat( + "ShardingParamShardingSpec has %d shards, but sharding param has " + "%d unique tiles, which is not a divisor of the number of shards", + num_shards_, num_unique_tiles)); + return; + } + const int replication = device_list.size() / num_unique_tiles; + + std::vector shard_to_unique_index_domain_index(num_shards_); + for (int i = 0; i < device_list.size(); ++i) { + const int device_idx = device_list[i]; + const int tile_idx = i / replication; + shard_to_unique_index_domain_index[device_idx] = tile_idx; + } + cached_shard_to_unique_index_domain_index_ = + std::move(shard_to_unique_index_domain_index); + }); + ABSL_RETURN_IF_ERROR(cached_shard_to_unique_index_domain_index_.status()); + return absl::MakeConstSpan(*cached_shard_to_unique_index_domain_index_); +} + std::string ShardingParamShardingSpec::DebugString() const { return absl::StrFormat("ShardingParamShardingSpec(num_shards: %d, %s)", num_shards_, sharding_param_.DebugString()); diff --git a/third_party/xla/xla/python/ifrt/sharding_spec.h b/third_party/xla/xla/python/ifrt/sharding_spec.h index bafdff441ba469..38b6941546ae24 100644 --- a/third_party/xla/xla/python/ifrt/sharding_spec.h +++ b/third_party/xla/xla/python/ifrt/sharding_spec.h @@ -24,11 +24,15 @@ limitations under the License. #include #include +#include "absl/base/attributes.h" +#include "absl/base/call_once.h" #include "absl/base/nullability.h" +#include "absl/container/inlined_vector.h" #include "absl/hash/hash.h" #include "absl/log/check.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" +#include "absl/types/span.h" #include "xla/python/ifrt/device_list.h" #include "xla/python/ifrt/index_domain.h" #include "xla/python/ifrt/ir/sharding_param.h" @@ -116,6 +120,40 @@ class ShardingSpec : public RTTIExtends, virtual absl::StatusOr> IndexDomains( const Shape& shape) const = 0; + struct IndexDomainAndShardIndices { + // The index domain mapped from shards. + IndexDomain index_domain; + + // All shard indices that map to this index domain. + absl::Span shard_indices; + + bool operator==(const IndexDomainAndShardIndices& other) const { + return index_domain == other.index_domain && + shard_indices == other.shard_indices; + } + bool operator!=(const IndexDomainAndShardIndices& other) const { + return !(*this == other); + } + }; + + // Breaks a shape up into unique `IndexDomain`s and the shard indices mapped + // to it. + // + // The result is valid for the lifetime of this `ShardingSpec`. + virtual absl::StatusOr> + UniqueIndexDomains(const Shape& shape) const + ABSL_ATTRIBUTE_LIFETIME_BOUND = 0; + + // Inverse of `UniqueIndexDomains()` for `shard_indices`. Does not take + // `shape` because the result is independent of `shape`. + // + // Suppose `j` be `unique_index_domain_indices[shard_i]`. Then, + // `unique_index_domains[j].shard_indices` contains `shard_i`. + // + // The result is valid for the lifetime of this `ShardingSpec`. + virtual absl::StatusOr> ShardToUniqueIndexDomainIndex() + const ABSL_ATTRIBUTE_LIFETIME_BOUND = 0; + template friend H AbslHashValue(H h, const ShardingSpec& value) { value.Hash(absl::HashState::Create(&h)); @@ -212,6 +250,12 @@ class SingleDeviceShardingSpec final absl::StatusOr> IndexDomains( const Shape& shape) const override; + absl::StatusOr> + UniqueIndexDomains(const Shape& shape) const override; + + absl::StatusOr> ShardToUniqueIndexDomainIndex() + const override; + static char ID; // NOLINT private: @@ -250,6 +294,12 @@ class OpaqueShardingSpec absl::StatusOr> IndexDomains( const Shape& shape) const override; + absl::StatusOr> + UniqueIndexDomains(const Shape& shape) const override; + + absl::StatusOr> ShardToUniqueIndexDomainIndex() + const override; + static char ID; // NOLINT private: @@ -340,8 +390,16 @@ class ConcreteShardingSpec absl::StatusOr> IndexDomains( const Shape& shape) const override; + absl::StatusOr> + UniqueIndexDomains(const Shape& shape) const override; + + absl::StatusOr> ShardToUniqueIndexDomainIndex() + const override; + static char ID; // NOLINT + ConcreteShardingSpec(const ConcreteShardingSpec& other); + private: ConcreteShardingSpec( int num_shards, Shape shape, std::vector shard_shapes, @@ -358,6 +416,12 @@ class ConcreteShardingSpec std::variant, std::vector> shard_shapes_; std::optional shard_shape_; std::optional> index_domains_; + + mutable absl::once_flag unique_shard_indices_once_; + mutable std::vector cached_shard_indices_; + mutable std::vector cached_shard_indices_offsets_; + mutable absl::once_flag shard_to_unique_index_domain_index_once_; + mutable std::vector cached_shard_to_unique_index_domain_index_; }; // Opaque sharding spec that does not define a fixed semantics for conversion @@ -402,8 +466,16 @@ class ConcreteEvenShardingSpec absl::StatusOr> IndexDomains( const Shape& shape) const override; + absl::StatusOr> + UniqueIndexDomains(const Shape& shape) const override; + + absl::StatusOr> ShardToUniqueIndexDomainIndex() + const override; + static char ID; // NOLINT + ConcreteEvenShardingSpec(const ConcreteEvenShardingSpec& other); + private: ConcreteEvenShardingSpec(int num_shards, Shape shape, Shape shard_shape, bool is_fully_replicated); @@ -414,6 +486,11 @@ class ConcreteEvenShardingSpec Shape shape_; Shape shard_shape_; + + mutable absl::once_flag unique_shard_indices_once_; + mutable std::vector cached_shard_indices_; + mutable absl::once_flag shard_to_unique_index_domain_index_once_; + mutable std::vector cached_shard_to_unique_index_domain_index_; }; // Sharding spec derived from an IR ShardingParam. @@ -441,8 +518,16 @@ class ShardingParamShardingSpec absl::StatusOr> IndexDomains( const Shape& shape) const override; + absl::StatusOr> + UniqueIndexDomains(const Shape& shape) const override; + + absl::StatusOr> ShardToUniqueIndexDomainIndex() + const override; + static char ID; // NOLINT + ShardingParamShardingSpec(const ShardingParamShardingSpec& other); + private: ShardingParamShardingSpec(int num_shards, ShardingParam sharding_param); @@ -451,6 +536,12 @@ class ShardingParamShardingSpec void Hash(absl::HashState state) const override; ShardingParam sharding_param_; + + mutable absl::once_flag unique_shard_indices_once_; + mutable absl::StatusOr> cached_shard_indices_; + mutable absl::once_flag shard_to_unique_index_domain_index_once_; + mutable absl::StatusOr> + cached_shard_to_unique_index_domain_index_; }; } // namespace ifrt diff --git a/third_party/xla/xla/python/ifrt/sharding_spec_test.cc b/third_party/xla/xla/python/ifrt/sharding_spec_test.cc index a1c7f2969b9338..86d45252180f5b 100644 --- a/third_party/xla/xla/python/ifrt/sharding_spec_test.cc +++ b/third_party/xla/xla/python/ifrt/sharding_spec_test.cc @@ -43,6 +43,7 @@ namespace { using ::absl_testing::StatusIs; using ::testing::ElementsAre; using ::testing::ElementsAreArray; +using ::testing::FieldsAre; using ::testing::HasSubstr; using ::testing::SizeIs; @@ -101,6 +102,17 @@ TEST_P(SingleDeviceShardingSpecTest, IndexDomains) { EXPECT_THAT(index_domains, ElementsAre(IndexDomain(shape))); } +TEST_P(SingleDeviceShardingSpecTest, UniqueIndexDomains) { + ShardingSpecRef sharding = SingleDeviceShardingSpec::Create(); + + Shape shape({10, 20}); + EXPECT_THAT(sharding->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds( + ElementsAre(FieldsAre(IndexDomain(shape), ElementsAre(0))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0))); +} + TEST_P(SingleDeviceShardingSpecTest, Disassemble) { ShardingSpecRef sharding = SingleDeviceShardingSpec::Create(); @@ -183,6 +195,21 @@ TEST_P(OpaqueShardingSpecTest, IndexDomainsFails) { "OpaqueShardingSpec does not have index domain information"))); } +TEST_P(OpaqueShardingSpecTest, UniqueIndexDomains) { + ShardingSpecRef sharding = OpaqueShardingSpec::Create(num_shards()); + + EXPECT_THAT( + sharding->UniqueIndexDomains(Shape({30})), + absl_testing::StatusIs( + tsl::error::INVALID_ARGUMENT, + HasSubstr("OpaqueShardingSpec does not support UniqueIndexDomains"))); + EXPECT_THAT( + sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("OpaqueShardingSpec does not support " + "ShardToUniqueIndexDomainIndex"))); +} + TEST_P(OpaqueShardingSpecTest, Hash) { EXPECT_TRUE(absl::VerifyTypeImplementsAbslHashCorrectly({ *OpaqueShardingSpec::Create(2), @@ -371,6 +398,44 @@ TEST_P(ConcreteShardingSpecTest, IndexDomainsMissing) { "ConcreteShardingSpec does not have index domain information"))); } +TEST_P(ConcreteShardingSpecTest, UniqueIndexDomains) { + std::vector shard_shapes{Shape({10}), Shape({10}), Shape({20})}; + std::vector index_domains{ + IndexDomain(Index({0}), Shape({10})), + IndexDomain(Index({0}), Shape({10})), + IndexDomain(Index({10}), Shape({20})), + }; + ShardingSpecRef sharding = + ConcreteShardingSpec::Create(Shape({30}), shard_shapes, index_domains); + + EXPECT_THAT( + sharding->UniqueIndexDomains(Shape({30})), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0}), Shape({10})), ElementsAre(0, 1)), + FieldsAre(IndexDomain(Index({10}), Shape({20})), ElementsAre(2))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 0, 1))); +} + +TEST_P(ConcreteShardingSpecTest, UniqueIndexDomainsMissing) { + std::vector shard_shapes{Shape({10}), Shape({20})}; + ShardingSpecRef sharding = + ConcreteShardingSpec::Create(Shape({30}), shard_shapes); + + EXPECT_THAT( + sharding->UniqueIndexDomains(Shape({30})), + absl_testing::StatusIs( + tsl::error::INVALID_ARGUMENT, + HasSubstr( + "ConcreteShardingSpec does not have index domain information"))); + EXPECT_THAT( + sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::StatusIs( + tsl::error::INVALID_ARGUMENT, + HasSubstr( + "ConcreteShardingSpec does not have index domain information"))); +} + TEST_P(ConcreteShardingSpecTest, Hash) { ASSERT_OK_AND_ASSIGN( auto dynamic_shape, @@ -505,6 +570,38 @@ TEST_P(ConcreteEvenShardingSpecTest, IndexDomainsFails) { "index domain information"))); } +TEST_P(ConcreteEvenShardingSpecTest, UniqueIndexDomains) { + { + // Fully replicated. + ShardingSpecRef sharding = ConcreteEvenShardingSpec::Create( + /*num_shards=*/2, Shape({30}), Shape({30}), + /*is_fully_replicated=*/true); + + EXPECT_THAT(sharding->UniqueIndexDomains(Shape({30})), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Shape({30})), ElementsAre(0, 1))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 0))); + } + { + // Not fully replicated. + ShardingSpecRef sharding = ConcreteEvenShardingSpec::Create( + /*num_shards=*/2, Shape({30}), Shape({15}), + /*is_fully_replicated=*/false); + + EXPECT_THAT(sharding->UniqueIndexDomains(Shape({30})), + absl_testing::StatusIs( + tsl::error::INVALID_ARGUMENT, + HasSubstr("ConcreteEvenShardingSpec does not have index " + "domain information"))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::StatusIs( + tsl::error::INVALID_ARGUMENT, + HasSubstr("ConcreteEvenShardingSpec does not have index " + "domain information"))); + } +} + TEST_P(ConcreteEvenShardingSpecTest, Hash) { EXPECT_TRUE(absl::VerifyTypeImplementsAbslHashCorrectly({ *ConcreteEvenShardingSpec::Create( @@ -679,6 +776,49 @@ TEST_P(ShardingParamShardingSpecTest, IndexDomainWithReplication) { } } +TEST_P(ShardingParamShardingSpecTest, UniqueIndexDomains) { + { + // 2x3 tiled with replication 1. + ShardingParam param{/*dim_shards=*/{2, 3}, + {/*permutation=*/{0, 1}, /*axis_sizes=*/{2, 3}}}; + ShardingSpecRef sharding = ShardingParamShardingSpec::Create(param); + + Shape shape({4, 9}); + EXPECT_THAT(sharding->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0, 0}), Shape({2, 3})), + ElementsAre(0)), + FieldsAre(IndexDomain(Index({0, 3}), Shape({2, 3})), + ElementsAre(1)), + FieldsAre(IndexDomain(Index({0, 6}), Shape({2, 3})), + ElementsAre(2)), + FieldsAre(IndexDomain(Index({2, 0}), Shape({2, 3})), + ElementsAre(3)), + FieldsAre(IndexDomain(Index({2, 3}), Shape({2, 3})), + ElementsAre(4)), + FieldsAre(IndexDomain(Index({2, 6}), Shape({2, 3})), + ElementsAre(5))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 1, 2, 3, 4, 5))); + } + { + // 2x1 tiled with replication 3. + ShardingParam param{/*dim_shards=*/{2, 1}, + {/*permutation=*/{0, 1}, /*axis_sizes=*/{2, 3}}}; + ShardingSpecRef sharding = ShardingParamShardingSpec::Create(param); + + Shape shape({4, 9}); + EXPECT_THAT(sharding->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0, 0}), Shape({2, 9})), + ElementsAre(0, 1, 2)), + FieldsAre(IndexDomain(Index({2, 0}), Shape({2, 9})), + ElementsAre(3, 4, 5))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 0, 0, 1, 1, 1))); + } +} + TEST_P(ShardingParamShardingSpecTest, IndexDomainZeroRank) { ShardingParam param{/*dim_shards=*/{}, {/*permutation=*/{0}, /*axis_sizes=*/{6}}}; diff --git a/third_party/xla/xla/python/ifrt/sharding_test.cc b/third_party/xla/xla/python/ifrt/sharding_test.cc index 9ddcdaffdc1b06..0c403771a16579 100644 --- a/third_party/xla/xla/python/ifrt/sharding_test.cc +++ b/third_party/xla/xla/python/ifrt/sharding_test.cc @@ -44,6 +44,7 @@ using ::absl_testing::IsOkAndHolds; using ::absl_testing::StatusIs; using ::testing::ElementsAre; using ::testing::ElementsAreArray; +using ::testing::FieldsAre; using ::testing::HasSubstr; using ::testing::SizeIs; @@ -142,6 +143,19 @@ TEST_P(SingleDeviceShardingTest, IndexDomains) { } } +TEST_P(SingleDeviceShardingTest, UniqueIndexDomains) { + DeviceListRef device_list = GetDevices({0}); + ShardingRef sharding = SingleDeviceSharding::Create( + device_list->devices().front(), MemoryKind()); + + Shape shape({10, 20}); + EXPECT_THAT( + sharding->UniqueIndexDomains(shape), + IsOkAndHolds(ElementsAre(FieldsAre(IndexDomain(shape), ElementsAre(0))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + IsOkAndHolds(ElementsAre(0))); +} + TEST_P(SingleDeviceShardingTest, Disassemble) { auto device_list = GetDevices({0}); ShardingRef sharding = SingleDeviceSharding::Create( @@ -298,6 +312,19 @@ TEST_P(OpaqueShardingTest, IndexDomainsFails) { HasSubstr("OpaqueSharding does not have index domain information"))); } +TEST_P(OpaqueShardingTest, UniqueIndexDomainsFails) { + DeviceListRef device_list = GetDevices({0, 1}); + ShardingRef sharding = OpaqueSharding::Create(device_list, MemoryKind()); + + EXPECT_THAT(sharding->UniqueIndexDomains(Shape({30})), + StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("does not support UniqueIndexDomains"))); + EXPECT_THAT( + sharding->ShardToUniqueIndexDomainIndex(), + StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("does not support ShardToUniqueIndexDomainIndex"))); +} + TEST_P(OpaqueShardingTest, Hash) { EXPECT_TRUE(absl::VerifyTypeImplementsAbslHashCorrectly({ *OpaqueSharding::Create(GetDevices({0, 1}), MemoryKind()), @@ -668,6 +695,35 @@ TEST_P(ConcreteShardingTest, IndexDomainsFails) { "of index domains and addressable devices"))); } +TEST_P(ConcreteShardingTest, UniqueIndexDomains) { + DeviceListRef device_list = GetDevices({0, 1, 2, 3, 4, 5}); + // devices 0..3 are addressable, 4..5 are non-addressable. + std::vector shard_shapes = { + Shape({10}), Shape({10}), Shape({10}), + Shape({10}), Shape({10}), Shape({10}), + }; + std::vector index_domains{ + IndexDomain(Index({0}), Shape({10})), + IndexDomain(Index({0}), Shape({10})), + IndexDomain(Index({10}), Shape({10})), + IndexDomain(Index({10}), Shape({10})), + IndexDomain(Index({20}), Shape({10})), + IndexDomain(Index({20}), Shape({10})), + }; + ShardingRef sharding = ConcreteSharding::Create( + device_list, MemoryKind(), Shape({30}), shard_shapes, index_domains); + + EXPECT_THAT( + sharding->UniqueIndexDomains(Shape({30})), + IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0}), Shape({10})), ElementsAre(0, 1)), + FieldsAre(IndexDomain(Index({10}), Shape({10})), ElementsAre(2, 3)), + FieldsAre(IndexDomain(Index({20}), Shape({10})), + ElementsAre(4, 5))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + IsOkAndHolds(ElementsAre(0, 0, 1, 1, 2, 2))); +} + TEST_P(ConcreteShardingTest, Hash) { ASSERT_OK_AND_ASSIGN( auto dynamic_shape, @@ -894,6 +950,38 @@ TEST_P(ConcreteEvenShardingTest, IndexDomainsFailsForNonFullyReplicated) { "ConcreteEvenSharding does not have index domain information"))); } +TEST_P(ConcreteEvenShardingTest, UniqueIndexDomains) { + Shape shape({10, 20}); + + auto device_list = GetDevices({0, 4}); + ASSERT_TRUE(device_list->devices()[0]->IsAddressable()); + ASSERT_FALSE(device_list->devices()[1]->IsAddressable()); + + ShardingRef sharding = ConcreteEvenSharding::Create( + device_list, MemoryKind(), /*shape=*/shape, /*shard_shape=*/shape, + /*is_fully_replicated=*/true); + + EXPECT_THAT(sharding->UniqueIndexDomains(shape), + IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(shape), ElementsAre(0, 1))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + IsOkAndHolds(ElementsAre(0, 0))); +} + +TEST_P(ConcreteEvenShardingTest, UniqueIndexDomainsFailsForNonFullyReplicated) { + auto device_list = GetDevices({0, 1}); + ShardingRef sharding = + ConcreteEvenSharding::Create(device_list, MemoryKind(), Shape({30}), + Shape({5}), /*is_fully_replicated=*/false); + + EXPECT_THAT(sharding->UniqueIndexDomains(Shape({30})), + StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("does not have index domain information"))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("does not have index domain information"))); +} + TEST_P(ConcreteEvenShardingTest, Hash) { EXPECT_TRUE(absl::VerifyTypeImplementsAbslHashCorrectly({ *ConcreteEvenSharding::Create(GetDevices({0, 1}), MemoryKind(), @@ -1237,6 +1325,27 @@ TEST_P(ShardingParamShardingTest, IndexDomainWithReplication) { } } +TEST_P(ShardingParamShardingTest, UniqueIndexDomains) { + auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); + // devices 0..3 are addressable, 4..5 are non-addressable. + // 2x1 tiled with replication 3. + ShardingParam param{/*dim_shards=*/{2, 1}, + {/*permutation=*/{0, 1}, /*axis_sizes=*/{2, 3}}}; + ASSERT_OK_AND_ASSIGN( + ShardingRef sharding, + ShardingParamSharding::Create(param, device_list, MemoryKind())); + + Shape shape({6, 6}); + EXPECT_THAT(sharding->UniqueIndexDomains(shape), + IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0, 0}), Shape({3, 6})), + ElementsAre(0, 1, 2)), + FieldsAre(IndexDomain(Index({3, 0}), Shape({3, 6})), + ElementsAre(3, 4, 5))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + IsOkAndHolds(ElementsAre(0, 0, 0, 1, 1, 1))); +} + TEST_P(ShardingParamShardingTest, IndexDomainZeroRank) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); ShardingParam param{/*dim_shards=*/{}, diff --git a/third_party/xla/xla/python/pjrt_ifrt/BUILD b/third_party/xla/xla/python/pjrt_ifrt/BUILD index abc818cfac4811..13bf8eb35a761c 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/BUILD +++ b/third_party/xla/xla/python/pjrt_ifrt/BUILD @@ -47,7 +47,9 @@ cc_library( "//xla/python/ifrt:serdes", "//xla/python/ifrt:serdes_version", "//xla/service:device_assignment", + "@com_google_absl//absl/base", "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/hash", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding.cc index 29091deb0a6a74..9bc5bc00053bd1 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding.cc @@ -26,6 +26,7 @@ limitations under the License. #include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" +#include "absl/status/statusor.h" #include "absl/strings/str_format.h" #include "absl/types/span.h" #include "xla/hlo/ir/hlo_sharding.h" diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding.h b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding.h index 38140cfa0ea491..13d26a44a3752b 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding.h +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding.h @@ -16,8 +16,6 @@ limitations under the License. #ifndef XLA_PYTHON_PJRT_IFRT_XLA_SHARDING_H_ #define XLA_PYTHON_PJRT_IFRT_XLA_SHARDING_H_ -#include -#include #include #include #include @@ -26,7 +24,6 @@ limitations under the License. #include "absl/hash/hash.h" #include "absl/status/statusor.h" -#include "absl/types/span.h" #include "xla/hlo/ir/hlo_sharding.h" #include "xla/python/ifrt/device_list.h" #include "xla/python/ifrt/index_domain.h" diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.cc index cc3a02453c152d..d92d8a643426c6 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.cc @@ -23,12 +23,15 @@ limitations under the License. #include #include +#include "absl/base/call_once.h" #include "absl/base/optimization.h" +#include "absl/container/inlined_vector.h" #include "absl/hash/hash.h" #include "absl/log/check.h" #include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" +#include "absl/status/statusor.h" #include "absl/strings/str_format.h" #include "absl/types/span.h" #include "xla/hlo/ir/hlo_sharding.h" @@ -40,6 +43,7 @@ limitations under the License. #include "xla/python/ifrt/shape.h" #include "xla/python/ifrt/sharding_spec.h" #include "xla/python/pjrt_ifrt/xla_sharding.h" +#include "xla/shape.h" #include "xla/shape_util.h" #include "xla/util.h" #include "xla/xla_data.pb.h" @@ -53,7 +57,6 @@ char HloShardingSpec::ID = 0; // NOLINT namespace { // Generates IndexDomains for an HloShardingSpec, using XLA HloSharding APIs. -// Note that this is O(N^2) where N is the number of devices (shards). std::vector IndexDomainsSlowPath( const xla::HloSharding& hlo_sharding, int num_shards, const Shape& shape) { // Only shape dimensions are used. @@ -108,6 +111,11 @@ HloShardingSpec::HloShardingSpec(int num_shards, num_shards_ == 1); } +HloShardingSpec::HloShardingSpec(const HloShardingSpec& other) + : RTTIExtends(other), + xla_hlo_sharding_(other.xla_hlo_sharding_), + hash_(other.hash_.load(std::memory_order_relaxed)) {} + absl::StatusOr HloShardingSpec::ToSharding( DeviceListRef devices, MemoryKind memory_kind) const { if (devices->size() != num_shards()) { @@ -281,6 +289,139 @@ absl::StatusOr> HloShardingSpec::IndexDomains( return result; } +absl::StatusOr> +HloShardingSpec::UniqueIndexDomains(const Shape& shape) const { + if (xla_hlo_sharding_.IsManual()) { + return absl::InvalidArgumentError( + "Manual sharding does not support UniqueIndexDomains"); + } + if (xla_hlo_sharding_.IsUnreduced()) { + return absl::InvalidArgumentError( + "Unreduced sharding does not support UniqueIndexDomains"); + } + if (xla_hlo_sharding_.HasNonReplicatedSubgroup()) { + return absl::InvalidArgumentError( + "Non-replicated subgroup (e.g., manual or unreduced subgroup) sharding " + "does not support UniqueIndexDomains"); + } + if (xla_hlo_sharding_.IsReplicatedOrSingleDevice()) { + absl::call_once(unique_shard_indices_once_, [this] { + cached_shard_indices_.reserve(num_shards_); + for (int i = 0; i < num_shards_; ++i) { + cached_shard_indices_.push_back(i); + } + }); + return absl::InlinedVector{ + IndexDomainAndShardIndices{ + /*index_domain=*/IndexDomain(shape), + /*shard_indices=*/absl::MakeConstSpan(cached_shard_indices_), + }, + }; + } + + const int64_t tiled_data_rank = xla_hlo_sharding_.TiledDataRank(); + if (shape.dims().size() != tiled_data_rank) { + return absl::InvalidArgumentError( + absl::StrFormat("shape must have %d dimensions, but has %d dimensions: " + "shape=%v, sharding=%s", + tiled_data_rank, shape.dims().size(), shape, + xla_hlo_sharding_.ToString())); + } + + absl::call_once(unique_shard_indices_once_, [this] { + const int64_t* flat_tile_assignment = + xla_hlo_sharding_.tile_assignment().array().data(); + cached_shard_indices_.reserve(num_shards_); + for (int64_t i = 0; i < num_shards_; ++i) { + cached_shard_indices_.push_back( + static_cast(flat_tile_assignment[i])); + } + }); + + const int64_t num_unique_tiles = xla_hlo_sharding_.NumTiles(); + if (num_shards_ % num_unique_tiles != 0) { + return absl::InvalidArgumentError(absl::StrFormat( + "HloShardingSpec has %d shards, but HloSharding has %d unique tiles, " + "which is not a divisor of the number of shards", + num_shards_, num_unique_tiles)); + } + const int64_t num_replicas = num_shards_ / num_unique_tiles; + + xla::Shape xla_shape = xla::ShapeUtil::MakeShapeWithDescendingLayout( + xla::PrimitiveType::S32, shape.dims()); + absl::InlinedVector unique_domains; + unique_domains.reserve(num_unique_tiles); + for (int64_t tile_idx = 0; tile_idx < num_unique_tiles; ++tile_idx) { + const int first_shard = cached_shard_indices_[tile_idx * num_replicas]; + std::vector tile_offset = + xla_hlo_sharding_.TileOffsetForDevice(xla_shape, first_shard); + std::vector tile_limit = + xla_hlo_sharding_.TileLimitForDevice(xla_shape, first_shard); + Index::Elements origin(shape.dims().size()); + Shape::Dimensions shard_shape(shape.dims().size()); + for (int i = 0; i < shape.dims().size(); ++i) { + origin[i] = tile_offset[i]; + shard_shape[i] = tile_limit[i] - tile_offset[i]; + } + unique_domains.push_back(IndexDomainAndShardIndices{ + /*index_domain=*/ + IndexDomain(Index(std::move(origin)), Shape(std::move(shard_shape))), + /*shard_indices=*/ + absl::MakeConstSpan(cached_shard_indices_) + .subspan(tile_idx * num_replicas, num_replicas), + }); + } + + return unique_domains; +} + +absl::StatusOr> +HloShardingSpec::ShardToUniqueIndexDomainIndex() const { + if (xla_hlo_sharding_.IsManual()) { + return absl::InvalidArgumentError( + "Manual sharding does not support ShardToUniqueIndexDomainIndex"); + } + if (xla_hlo_sharding_.IsUnreduced()) { + return absl::InvalidArgumentError( + "Unreduced sharding does not support ShardToUniqueIndexDomainIndex"); + } + if (xla_hlo_sharding_.HasNonReplicatedSubgroup()) { + return absl::InvalidArgumentError( + "Non-replicated subgroup (e.g., manual or unreduced subgroup) sharding " + "does not support ShardToUniqueIndexDomainIndex"); + } + if (xla_hlo_sharding_.IsReplicatedOrSingleDevice()) { + absl::call_once(shard_to_unique_index_domain_index_once_, [this] { + cached_shard_to_unique_index_domain_index_.assign(num_shards_, 0); + }); + return absl::MakeConstSpan(cached_shard_to_unique_index_domain_index_); + } + + const int64_t num_unique_tiles = xla_hlo_sharding_.NumTiles(); + if (num_shards_ % num_unique_tiles != 0) { + return absl::InvalidArgumentError(absl::StrFormat( + "HloShardingSpec has %d shards, but HloSharding has %d unique tiles, " + "which is not a divisor of the number of shards", + num_shards_, num_unique_tiles)); + } + const int64_t num_replicas = num_shards_ / num_unique_tiles; + + absl::call_once(shard_to_unique_index_domain_index_once_, [&, this] { + cached_shard_to_unique_index_domain_index_.resize(num_shards_); + const int64_t* flat_tile_assignment = + xla_hlo_sharding_.tile_assignment().array().data(); + for (int64_t tile_idx = 0; tile_idx < num_unique_tiles; ++tile_idx) { + const int64_t offset = tile_idx * num_replicas; + for (int64_t i = 0; i < num_replicas; ++i) { + const int device_idx = + static_cast(flat_tile_assignment[offset + i]); + cached_shard_to_unique_index_domain_index_[device_idx] = tile_idx; + } + } + }); + return absl::MakeConstSpan(cached_shard_to_unique_index_domain_index_); +} + std::string HloShardingSpec::DebugString() const { return absl::StrFormat("HloShardingSpec(num_shards: %d, hlo_sharding: %s)", num_shards_, xla_hlo_sharding_.ToString()); diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.h b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.h index b2685077274f3a..4d5fc96ad59d53 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.h +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.h @@ -23,6 +23,7 @@ limitations under the License. #include #include +#include "absl/base/call_once.h" #include "absl/hash/hash.h" #include "absl/status/statusor.h" #include "xla/hlo/ir/hlo_sharding.h" @@ -79,8 +80,16 @@ class HloShardingSpec final absl::StatusOr> IndexDomains( const Shape& shape) const override; + absl::StatusOr> + UniqueIndexDomains(const Shape& shape) const override; + + absl::StatusOr> ShardToUniqueIndexDomainIndex() + const override; + static char ID; // NOLINT + HloShardingSpec(const HloShardingSpec& other); + private: HloShardingSpec(int num_shards, xla::HloSharding xla_hlo_sharding); @@ -94,6 +103,11 @@ class HloShardingSpec final // May be written multiple times with the same non-zero value. static constexpr uint64_t kUnsetHash = 0; mutable std::atomic hash_ = kUnsetHash; + + mutable absl::once_flag unique_shard_indices_once_; + mutable std::vector cached_shard_indices_; + mutable absl::once_flag shard_to_unique_index_domain_index_once_; + mutable std::vector cached_shard_to_unique_index_domain_index_; }; // Test only: returns `HloShardingSpec::IndexDomains()`, using diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_test.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_test.cc index 99cc2e0411c5c2..214bff8445f9e3 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_test.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_test.cc @@ -42,6 +42,7 @@ namespace { using ::testing::ElementsAre; using ::testing::ElementsAreArray; +using ::testing::FieldsAre; using ::testing::HasSubstr; using ::testing::SizeIs; @@ -206,6 +207,21 @@ TEST_F(HloShardingSpecTest, IndexDomainsWithReplication) { ElementsAreArray(TEST_HloShardingSpecIndexDomainsSlowPath(*spec, shape))); } +TEST_F(HloShardingSpecTest, UniqueIndexDomainsWithReplication) { + int num_shards = 6; + // Fully replicated. + auto xla_hlo_sharding = xla::HloSharding::Replicate(); + std::shared_ptr spec = + HloShardingSpec::Create(num_shards, xla_hlo_sharding); + + Shape shape({10, 20}); + EXPECT_THAT(spec->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre(FieldsAre( + IndexDomain(shape), ElementsAre(0, 1, 2, 3, 4, 5))))); + EXPECT_THAT(spec->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 0, 0, 0, 0, 0))); +} + TEST_F(HloShardingSpecTest, DisassembleWithReplication) { int num_shards = 6; // Fully replicated. @@ -244,6 +260,28 @@ TEST_F(HloShardingSpecTest, IndexDomainsWithTile) { ElementsAreArray(TEST_HloShardingSpecIndexDomainsSlowPath(*spec, shape))); } +TEST_F(HloShardingSpecTest, UniqueIndexDomainsWithTile) { + int num_shards = 6; + // 6-way sharded along axis 0, 1-way sharded along axis 1. + auto xla_hlo_sharding = xla::HloSharding::Tile(xla::TileAssignment({6, 1})); + std::shared_ptr spec = + HloShardingSpec::Create(num_shards, xla_hlo_sharding); + + Shape shape({12, 20}); + EXPECT_THAT( + spec->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0, 0}), Shape({2, 20})), ElementsAre(0)), + FieldsAre(IndexDomain(Index({2, 0}), Shape({2, 20})), ElementsAre(1)), + FieldsAre(IndexDomain(Index({4, 0}), Shape({2, 20})), ElementsAre(2)), + FieldsAre(IndexDomain(Index({6, 0}), Shape({2, 20})), ElementsAre(3)), + FieldsAre(IndexDomain(Index({8, 0}), Shape({2, 20})), ElementsAre(4)), + FieldsAre(IndexDomain(Index({10, 0}), Shape({2, 20})), + ElementsAre(5))))); + EXPECT_THAT(spec->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 1, 2, 3, 4, 5))); +} + TEST_F(HloShardingSpecTest, DisassembleWithTile) { int num_shards = 6; // 6-way sharded along axis 0, 1-way sharded along axis 1. @@ -282,6 +320,28 @@ TEST_F(HloShardingSpecTest, IndexDomainsWithUnevenTile) { ElementsAreArray(TEST_HloShardingSpecIndexDomainsSlowPath(*spec, shape))); } +TEST_F(HloShardingSpecTest, UniqueIndexDomainsWithUnevenTile) { + int num_shards = 6; + // 6-way sharded along axis 0, 1-way sharded along axis 1. + auto xla_hlo_sharding = xla::HloSharding::Tile(xla::TileAssignment({6, 1})); + std::shared_ptr spec = + HloShardingSpec::Create(num_shards, xla_hlo_sharding); + + Shape shape({11, 20}); + EXPECT_THAT( + spec->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0, 0}), Shape({2, 20})), ElementsAre(0)), + FieldsAre(IndexDomain(Index({2, 0}), Shape({2, 20})), ElementsAre(1)), + FieldsAre(IndexDomain(Index({4, 0}), Shape({2, 20})), ElementsAre(2)), + FieldsAre(IndexDomain(Index({6, 0}), Shape({2, 20})), ElementsAre(3)), + FieldsAre(IndexDomain(Index({8, 0}), Shape({2, 20})), ElementsAre(4)), + FieldsAre(IndexDomain(Index({10, 0}), Shape({1, 20})), + ElementsAre(5))))); + EXPECT_THAT(spec->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 1, 2, 3, 4, 5))); +} + TEST_F(HloShardingSpecTest, DisassembleWithUnevenTile) { int num_shards = 6; // 6-way sharded along axis 0, 1-way sharded along axis 1. @@ -326,6 +386,45 @@ TEST_F(HloShardingSpecTest, IndexDomainsWithPartialTile) { ElementsAreArray(TEST_HloShardingSpecIndexDomainsSlowPath(*spec, shape))); } +TEST_F(HloShardingSpecTest, UniqueIndexDomainsWithPartialTile) { + int num_shards = 6; + // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard + // replicated by 3 times. + auto xla_hlo_sharding = + xla::HloSharding::PartialTile(xla::TileAssignment({2, 1, 3})); + std::shared_ptr spec = + HloShardingSpec::Create(num_shards, xla_hlo_sharding); + + Shape shape({10, 20}); + EXPECT_THAT(spec->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0, 0}), Shape({5, 20})), + ElementsAre(0, 1, 2)), + FieldsAre(IndexDomain(Index({5, 0}), Shape({5, 20})), + ElementsAre(3, 4, 5))))); + EXPECT_THAT(spec->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 0, 0, 1, 1, 1))); +} + +TEST_F(HloShardingSpecTest, UniqueIndexDomainsManualFails) { + int num_shards = 6; + auto xla_hlo_sharding = xla::HloSharding::Manual(); + std::shared_ptr spec = + HloShardingSpec::Create(num_shards, xla_hlo_sharding); + + Shape shape({10, 20}); + EXPECT_THAT( + spec->UniqueIndexDomains(shape), + absl_testing::StatusIs( + tsl::error::INVALID_ARGUMENT, + HasSubstr("Manual sharding does not support UniqueIndexDomains"))); + EXPECT_THAT( + spec->ShardToUniqueIndexDomainIndex(), + absl_testing::StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("Manual sharding does not support " + "ShardToUniqueIndexDomainIndex"))); +} + TEST_F(HloShardingSpecTest, DisassembleWithPartialTile) { int num_shards = 6; // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_test.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_test.cc index ca24cd1d6aa397..d22c8afd621538 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_test.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_test.cc @@ -46,6 +46,7 @@ namespace { using ::testing::ElementsAre; using ::testing::ElementsAreArray; +using ::testing::FieldsAre; using ::testing::HasSubstr; using ::testing::SizeIs; @@ -277,6 +278,21 @@ TEST_P(HloShardingTest, IndexDomainsWithReplication) { } } +TEST_P(HloShardingTest, UniqueIndexDomainsWithReplication) { + auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); + // Fully replicated. + auto xla_hlo_sharding = xla::HloSharding::Replicate(); + std::shared_ptr sharding = + HloSharding::Create(device_list, MemoryKind(), xla_hlo_sharding); + + Shape shape({10, 20}); + EXPECT_THAT(sharding->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre(FieldsAre( + IndexDomain(shape), ElementsAre(0, 1, 2, 3, 4, 5))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 0, 0, 0, 0, 0))); +} + TEST_P(HloShardingTest, DisassembleWithReplication) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); // Fully replicated. @@ -354,6 +370,28 @@ TEST_P(HloShardingTest, IndexDomainsWithTile) { } } +TEST_P(HloShardingTest, UniqueIndexDomainsWithTile) { + auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); + // 6-way sharded along axis 0, 1-way sharded along axis 1. + auto xla_hlo_sharding = xla::HloSharding::Tile(xla::TileAssignment({6, 1})); + std::shared_ptr sharding = + HloSharding::Create(device_list, MemoryKind(), xla_hlo_sharding); + + Shape shape({12, 20}); + EXPECT_THAT( + sharding->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0, 0}), Shape({2, 20})), ElementsAre(0)), + FieldsAre(IndexDomain(Index({2, 0}), Shape({2, 20})), ElementsAre(1)), + FieldsAre(IndexDomain(Index({4, 0}), Shape({2, 20})), ElementsAre(2)), + FieldsAre(IndexDomain(Index({6, 0}), Shape({2, 20})), ElementsAre(3)), + FieldsAre(IndexDomain(Index({8, 0}), Shape({2, 20})), ElementsAre(4)), + FieldsAre(IndexDomain(Index({10, 0}), Shape({2, 20})), + ElementsAre(5))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 1, 2, 3, 4, 5))); +} + TEST_P(HloShardingTest, DisassembleWithTile) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); // 6-way sharded along axis 0, 1-way sharded along axis 1. @@ -474,7 +512,7 @@ TEST_P(HloShardingTest, DisassembleWithUnevenTile) { TEST_P(HloShardingTest, IndexDomainsWithPartialTile) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard - // replicated by 3 times. + // replicated 3 times. auto xla_hlo_sharding = xla::HloSharding::PartialTile(xla::TileAssignment({2, 1, 3})); std::shared_ptr sharding = @@ -514,10 +552,30 @@ TEST_P(HloShardingTest, IndexDomainsWithPartialTile) { } } +TEST_P(HloShardingTest, UniqueIndexDomainsWithPartialTile) { + auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); + // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard + // replicated 3 times. + auto xla_hlo_sharding = + xla::HloSharding::PartialTile(xla::TileAssignment({2, 1, 3})); + std::shared_ptr sharding = + HloSharding::Create(device_list, MemoryKind(), xla_hlo_sharding); + + Shape shape({10, 20}); + EXPECT_THAT(sharding->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0, 0}), Shape({5, 20})), + ElementsAre(0, 1, 2)), + FieldsAre(IndexDomain(Index({5, 0}), Shape({5, 20})), + ElementsAre(3, 4, 5))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 0, 0, 1, 1, 1))); +} + TEST_P(HloShardingTest, DisassembleWithPartialTile) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard - // replicated by 3 times. + // replicated 3 times. auto xla_hlo_sharding = xla::HloSharding::PartialTile(xla::TileAssignment({2, 1, 3})); std::shared_ptr sharding = @@ -555,7 +613,7 @@ TEST_P(HloShardingTest, DisassembleWithPartialTile) { TEST_P(HloShardingTest, IndexDomainsWithSubgroupReplicated) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard - // replicated by 3 times. + // replicated 3 times. auto xla_hlo_sharding = xla::HloSharding::Subgroup( xla::TileAssignment({2, 1, 3}), {xla::OpSharding::REPLICATED}); std::shared_ptr sharding = @@ -598,7 +656,7 @@ TEST_P(HloShardingTest, IndexDomainsWithSubgroupReplicated) { TEST_P(HloShardingTest, DisassembleWithSubgroupReplicated) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard - // replicated by 3 times. + // replicated 3 times. auto xla_hlo_sharding = xla::HloSharding::Subgroup( xla::TileAssignment({2, 1, 3}), {xla::OpSharding::REPLICATED}); std::shared_ptr sharding = @@ -636,7 +694,7 @@ TEST_P(HloShardingTest, DisassembleWithSubgroupReplicated) { TEST_P(HloShardingTest, IndexDomainsWithSubgroupMaximalSlowPath) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard - // maximal-replicated by 3 times, device#0 in each replication is maximal. + // maximal-replicated 3 times, device#0 in each replication is maximal. auto xla_hlo_sharding = xla::HloSharding::Subgroup( xla::TileAssignment({2, 1, 3}), {xla::OpSharding::MAXIMAL}); std::shared_ptr sharding = @@ -679,7 +737,7 @@ TEST_P(HloShardingTest, IndexDomainsWithSubgroupMaximalSlowPath) { TEST_P(HloShardingTest, DisassembleWithSubgroupMaximalSlowPath) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard - // maximal-replicated by 3 times, device#0 in each replication is maximal. + // maximal-replicated 3 times, device#0 in each replication is maximal. auto xla_hlo_sharding = xla::HloSharding::Subgroup( xla::TileAssignment({2, 1, 3}), {xla::OpSharding::MAXIMAL}); std::shared_ptr sharding = From afbf59b50c9a696c7721bbbf7cfc2e14382d8269 Mon Sep 17 00:00:00 2001 From: Zac Mustin Date: Wed, 2 Sep 2026 17:44:56 -0700 Subject: [PATCH 10/15] Implement `PjRtCApiClient::Compile` for `XlaComputation`. Implements `PjRtCApiClient::Compile(const XlaComputation&, CompileOptions)` so that clients can perform AoT compilation w/ the PJRT C API without immediately loading to device memory. PiperOrigin-RevId: 975423410 --- .../xla/pjrt/c_api_client/pjrt_c_api_client.cc | 9 +++++++++ .../xla/pjrt/c_api_client/pjrt_c_api_client.h | 3 +++ .../pjrt/c_api_client/pjrt_c_api_client_test.cc | 17 +++++++++++++++++ 3 files changed, 29 insertions(+) diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc index 9e08a8071358bb..7ea1ed00f02ba1 100644 --- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc +++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc @@ -747,6 +747,15 @@ InitializeArgsAndCompileAot(const PJRT_Api* c_api, PjRtClient* client, } // namespace +absl::StatusOr> PjRtCApiClient::Compile( + const XlaComputation& computation, CompileOptions options) { + tsl::profiler::TraceMe traceme("PjRtCApiClient::Compile(XlaComputation)"); + ABSL_ASSIGN_OR_RETURN(const PjRtTopologyDescription* const topology, + GetTopologyDescription()); + return InitializeArgsAndCompileAot(c_api_, this, &computation, options, + *topology); +} + absl::StatusOr> PjRtCApiClient::Compile( MaybeOwningMlirModule module, CompileOptions options) { tsl::profiler::TraceMe traceme([&module]() { diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h index e9fdfd162fc88f..4b4da58f2ffdcc 100644 --- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h +++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h @@ -422,6 +422,9 @@ class PjRtCApiClient : public PjRtClient { absl::StatusOr GetDefaultLayout( PrimitiveType element_type, absl::Span dims) override; + absl::StatusOr> Compile( + const XlaComputation& computation, CompileOptions options) override; + absl::StatusOr> CompileAndLoad( const XlaComputation& computation, CompileOptions options) override; diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc index 3c47f0ea8e8a07..b7efe0487da9d0 100644 --- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc +++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc @@ -573,6 +573,23 @@ TEST(PjRtClientTest, CompileMlirModule) { EXPECT_NE(executable.get(), nullptr); } +TEST(PjRtClientTest, CompileXlaComputation) { + SetUpCpuPjRtApi(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr client, + GetCApiClient("cpu")); + Shape shape = ShapeUtil::MakeShape(S32, {4}); + XlaBuilder builder("add_one"); + auto input = Parameter(&builder, 0, shape, "input"); + auto one = ConstantR0(&builder, 1); + auto add = Add(input, one); + ASSERT_OK_AND_ASSIGN(XlaComputation computation, builder.Build(add)); + + CompileOptions options; + ASSERT_OK_AND_ASSIGN(std::unique_ptr executable, + client->Compile(computation, options)); + EXPECT_NE(executable.get(), nullptr); +} + TEST(PjRtCApiClientTest, LoadExecutable) { SetUpCpuPjRtApi(); ASSERT_OK_AND_ASSIGN(std::unique_ptr client, From b2323ca068870d24daaef18fb0e26215e0854d44 Mon Sep 17 00:00:00 2001 From: Jian Cai Date: Wed, 2 Sep 2026 18:39:20 -0700 Subject: [PATCH 11/15] Update original value in AlgebraicSimplifier when pruning optimization barrier [XLA][HLO Value Tracking] Update original value in AlgebraicSimplifier when pruning optimization barrier When AlgebraicSimplifier prunes unused elements from an optimization barrier, update the instruction original value for both the barrier and its new operand using CopyOriginalValue from hlo_original_value_util.h. Also add vector support to CopyOriginalValue for vector-based index mappings. PiperOrigin-RevId: 975441821 --- .../xla/xla/hlo/ir/hlo_original_value_test.cc | 29 ++++++++++++++++++ .../xla/xla/hlo/ir/hlo_original_value_util.h | 28 ++++++++++++++++- .../simplifiers/algebraic_simplifier.cc | 3 ++ .../simplifiers/algebraic_simplifier_test.cc | 30 +++++++++++++++++++ 4 files changed, 89 insertions(+), 1 deletion(-) diff --git a/third_party/xla/xla/hlo/ir/hlo_original_value_test.cc b/third_party/xla/xla/hlo/ir/hlo_original_value_test.cc index 3ebc8b62570638..6b8553aa87af37 100644 --- a/third_party/xla/xla/hlo/ir/hlo_original_value_test.cc +++ b/third_party/xla/xla/hlo/ir/hlo_original_value_test.cc @@ -500,5 +500,34 @@ TEST_F(OriginalValueHloTest, CopyOriginalValueWithMap) { Optional(Eq(OriginalArray{"instA", {0}}))); } +TEST_F(OriginalValueHloTest, CopyOriginalValueWithVector) { + const char* hlo_string = R"( +HloModule test + +ENTRY main { + p0 = f32[] parameter(0) + p1 = f32[] parameter(1) + p2 = f32[] parameter(2) + ROOT tuple = (f32[], f32[], f32[]) tuple(p0, p1, p2), origin={({"p0"}, {"p1"}, {"p2"})} +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloInstruction* root = module->entry_computation()->root_instruction(); + HloInstruction* dest = + module->entry_computation()->AddInstruction(HloInstruction::CreateTuple( + {root->mutable_operand(2), root->mutable_operand(0)})); + + // Index 0 -> 1, Index 1 is pruned (-1), Index 2 -> 0. + std::vector old_to_new_tuple_idx = {1, -1, 0}; + + CopyOriginalValue(root, dest, old_to_new_tuple_idx); + + ASSERT_NE(dest->original_value(), nullptr); + EXPECT_THAT(dest->original_value()->original_array({0}), + Optional(Eq(OriginalArray{"p2"}))); + EXPECT_THAT(dest->original_value()->original_array({1}), + Optional(Eq(OriginalArray{"p0"}))); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/hlo/ir/hlo_original_value_util.h b/third_party/xla/xla/hlo/ir/hlo_original_value_util.h index c3e4740e538176..671127ce011072 100644 --- a/third_party/xla/xla/hlo/ir/hlo_original_value_util.h +++ b/third_party/xla/xla/hlo/ir/hlo_original_value_util.h @@ -16,9 +16,11 @@ limitations under the License. #ifndef XLA_HLO_IR_HLO_ORIGINAL_VALUE_UTIL_H_ #define XLA_HLO_IR_HLO_ORIGINAL_VALUE_UTIL_H_ +#include #include #include #include +#include #include "absl/container/flat_hash_map.h" #include "xla/hlo/ir/hlo_instruction.h" @@ -28,8 +30,12 @@ limitations under the License. namespace xla { // Checks if the type of the map is a matching integer map. +template +struct is_matching_integer_map : std::false_type {}; + template -struct is_matching_integer_map { +struct is_matching_integer_map< + T, std::void_t> { static constexpr bool value = std::is_integral::value && std::is_same::value; @@ -81,6 +87,26 @@ CopyOriginalValue(const HloInstruction* src_instruction, dest_instruction->set_original_value(new_original_value); } +// Copies the original value of the source to the destination instruction. +// Original arrays in the source original value are rearranged in the new +// original value according to the given vector of old to new tuple indices. +// Elements with negative values in the vector are treated as pruned/unused. +template +std::enable_if_t> CopyOriginalValue( + const HloInstruction* src_instruction, HloInstruction* dest_instruction, + const std::vector& old_to_new_tuple_idx) { + absl::flat_hash_map mapping; + for (size_t old_idx = 0; old_idx < old_to_new_tuple_idx.size(); ++old_idx) { + if constexpr (std::is_signed_v) { + if (old_to_new_tuple_idx[old_idx] < 0) { + continue; + } + } + mapping[static_cast(old_idx)] = old_to_new_tuple_idx[old_idx]; + } + CopyOriginalValue(src_instruction, dest_instruction, mapping); +} + // Copies the original value of the source to the destination instruction if the // shapes of the source and destination are compatible. This performs a deep // copy if clone is set to true. Otherwise, it performs a shallow copy. Print a diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc index a72a150261dba4..2a3ef2c1cb5aef 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc @@ -54,6 +54,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_original_value.h" +#include "xla/hlo/ir/hlo_original_value_util.h" #include "xla/hlo/ir/hlo_sharding.h" #include "xla/hlo/transforms/simplifiers/conv_operand_swapper.h" #include "xla/hlo/utils/hlo_sharding_util.h" @@ -5430,6 +5431,8 @@ absl::Status AlgebraicSimplifierVisitor::HandleOptimizationBarrier( operand->AddInstruction(HloInstruction::CreateTuple(operands)); ABSL_RETURN_IF_ERROR(barrier->ReplaceOperandWithDifferentShape(0, new_operand)); *barrier->mutable_shape() = new_operand->shape(); + CopyOriginalValue(barrier, barrier, index_map); + CopyOriginalValue(operand, new_operand, index_map); for (auto use : barrier->users()) { CHECK_EQ(use->opcode(), HloOpcode::kGetTupleElement); use->set_tuple_index(index_map[use->tuple_index()]); diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc index 5c827aec5f682e..3d80a30755bf1a 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc @@ -12669,6 +12669,36 @@ TEST_F(AlgebraicSimplifierTest, SimplifyOptimizationBarrier) { 2); } +TEST_F(AlgebraicSimplifierTest, SimplifyOptimizationBarrierWithOriginalValue) { + constexpr absl::string_view kModuleStr = R"( + HloModule m + + ENTRY entry { + param.0 = f32[] parameter(0) + param.1 = f32[] parameter(1) + add.0 = f32[] add(param.0, param.1) + sub.0 = f32[] subtract(param.0, param.1) + mul.0 = f32[] multiply(param.0, param.1) + tuple.0 = (f32[], f32[], f32[]) tuple(mul.0, sub.0, add.0), origin={({"mul.0"}, {"sub.0"}, {"add.0"})} + b = (f32[], f32[], f32[]) opt-barrier(tuple.0), origin={({"b" {0}}, {"b" {1}}, {"b" {2}})} + gte.0 = f32[] get-tuple-element(b), index=1 + ROOT t = (f32[]) tuple(gte.0) + } + )"; + ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + ASSERT_OK_AND_ASSIGN(bool changed, + AlgebraicSimplifier(default_options_).Run(m.get())); + EXPECT_TRUE(changed); + ASSERT_THAT(verifier().Run(m.get()), absl_testing::IsOk()); + const HloInstruction* b = FindInstruction(m.get(), "b"); + ASSERT_NE(b, nullptr); + EXPECT_EQ(b->shape().tuple_shapes().size(), 1); + ASSERT_NE(b->original_value(), nullptr); + EXPECT_EQ(b->original_value()->original_array({0})->instruction_name, "b"); + EXPECT_EQ(b->original_value()->original_array({0})->shape_index, + ShapeIndex({1})); +} + TEST_F(AlgebraicSimplifierTest, DoNotSimplifyOptimizationBarrierSideEffects) { constexpr absl::string_view kModuleStr = R"( HloModule m From 757d906c938a51785e7729657e3b99bbde2d830d Mon Sep 17 00:00:00 2001 From: James Spooner Date: Wed, 2 Sep 2026 18:44:42 -0700 Subject: [PATCH 12/15] - Add HloBuffer::ComputeSize(size_fn) and HloBuffer::IsHeapPressureImpacting(alloc_constants, buffers_to_assign). - Add schedule-independent HloLiveRange liveness helpers (GetBuffersDefined, GetBuffersUsed, GetBytesDefined, GetParameterBytesAtStart, BufferLivesOut). - Add unit tests in hlo_live_range_test. PiperOrigin-RevId: 975443269 --- third_party/xla/xla/hlo/utils/BUILD | 3 + .../xla/xla/hlo/utils/hlo_live_range.cc | 163 ++++++++++++++---- .../xla/xla/hlo/utils/hlo_live_range.h | 44 +++++ .../xla/xla/hlo/utils/hlo_live_range_test.cc | 62 +++++++ third_party/xla/xla/service/BUILD | 2 + .../xla/xla/service/buffer_assignment.h | 6 +- .../service/heap_simulator/heap_simulator.cc | 23 +-- .../service/heap_simulator/heap_simulator.h | 5 +- third_party/xla/xla/service/hlo_buffer.cc | 34 ++++ third_party/xla/xla/service/hlo_buffer.h | 20 +++ 10 files changed, 310 insertions(+), 52 deletions(-) diff --git a/third_party/xla/xla/hlo/utils/BUILD b/third_party/xla/xla/hlo/utils/BUILD index c5aa2126276acc..c4b0ff8a532a9c 100644 --- a/third_party/xla/xla/hlo/utils/BUILD +++ b/third_party/xla/xla/hlo/utils/BUILD @@ -78,8 +78,10 @@ cc_library( ":hlo_stack_trace", "//xla:shape_util", "//xla/hlo/analysis:hlo_alias_analysis", + "//xla/hlo/analysis:hlo_dataflow_analysis", "//xla/hlo/ir:hlo", "//xla/hlo/ir:hlo_instruction_utils", + "//xla/service:buffer_value", "//xla/service:hlo_buffer", "//xla/service:hlo_value", "@com_google_absl//absl/algorithm:container", @@ -110,6 +112,7 @@ xla_cc_test( "//xla/hlo/ir:hlo", "//xla/hlo/parser:hlo_parser", "//xla/hlo/testlib:hlo_hardware_independent_test_base", + "//xla/service:buffer_value", "//xla/service:hlo_value", "//xla/tsl/lib/core:status_test_util", "//xla/tsl/platform:statusor", diff --git a/third_party/xla/xla/hlo/utils/hlo_live_range.cc b/third_party/xla/xla/hlo/utils/hlo_live_range.cc index 6288d09d8723c5..1db60423a53b3e 100644 --- a/third_party/xla/xla/hlo/utils/hlo_live_range.cc +++ b/third_party/xla/xla/hlo/utils/hlo_live_range.cc @@ -37,12 +37,14 @@ limitations under the License. #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "xla/hlo/analysis/hlo_alias_analysis.h" +#include "xla/hlo/analysis/hlo_dataflow_analysis.h" #include "xla/hlo/ir/dfs_hlo_visitor.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction_utils.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_schedule.h" #include "xla/hlo/utils/hlo_stack_trace.h" +#include "xla/service/buffer_value.h" #include "xla/service/hlo_buffer.h" #include "xla/service/hlo_value.h" #include "xla/shape.h" @@ -63,6 +65,119 @@ absl::StatusOr> HloLiveRange::Run( return hlo_live_range; } +/*static*/ +std::vector HloLiveRange::GetValuesDefined( + const HloInstruction* instruction, const HloDataflowAnalysis& dataflow) { + std::vector values; + const auto& value_set_tree = dataflow.GetInstructionValueSet(instruction); + for (const auto& entry : value_set_tree) { + if (dataflow.ValueIsDefinedAt(instruction, entry.first)) { + values.push_back(&dataflow.GetValueDefinedAt(instruction, entry.first)); + } + } + return values; +} + +/*static*/ +std::vector HloLiveRange::GetBuffersDefined( + const HloInstruction* instruction, const HloAliasAnalysis& alias_analysis) { + std::vector buffers; + absl::flat_hash_set seen; + for (const HloValue* value : + GetValuesDefined(instruction, alias_analysis.dataflow_analysis())) { + const HloBuffer* buffer = &alias_analysis.GetBufferContainingValue(*value); + if (seen.insert(buffer).second) { + buffers.push_back(buffer); + } + } + return buffers; +} + +/*static*/ +int64_t HloLiveRange::GetBytesDefined( + const HloInstruction* instruction, const HloAliasAnalysis& alias_analysis, + const BufferValue::SizeFunction& size_fn) { + if (instruction->opcode() == HloOpcode::kParameter) { + return 0; + } + int64_t bytes = 0; + for (const HloBuffer* buffer : + GetBuffersDefined(instruction, alias_analysis)) { + if (buffer->IsHeapPressureImpacting()) { + bytes += buffer->ComputeSize(size_fn); + } + } + return bytes; +} + +/*static*/ +std::vector HloLiveRange::GetBuffersUsed( + const HloInstruction* instruction, const HloAliasAnalysis& alias_analysis) { + std::vector buffers; + absl::flat_hash_set seen; + const auto& dataflow = alias_analysis.dataflow_analysis(); + for (const HloInstruction* operand : instruction->operands()) { + HloValueSet value_set = dataflow.GetFlattenedValueSet(operand); + for (const HloValue* value : value_set.values()) { + const HloBuffer* buffer = + &alias_analysis.GetBufferContainingValue(*value); + if (seen.insert(buffer).second) { + buffers.push_back(buffer); + } + } + } + return buffers; +} + +/*static*/ +int32_t HloLiveRange::GetTotalUsers(const HloBuffer& buffer, + const HloComputation* computation) { + int32_t total = 0; + for (const HloValue* value : buffer.values()) { + for (const HloUse& use : value->GetUses()) { + if (computation == nullptr || use.instruction->parent() == computation) { + ++total; + } + } + } + return total; +} + +/*static*/ +int64_t HloLiveRange::GetParameterBytesAtStart( + const HloComputation& computation, const HloAliasAnalysis& alias_analysis, + const BufferValue::SizeFunction& size_fn) { + int64_t bytes = 0; + absl::flat_hash_set seen; + for (const HloInstruction* param : computation.parameter_instructions()) { + for (const HloBuffer* buffer : GetBuffersDefined(param, alias_analysis)) { + if (seen.insert(buffer).second && buffer->IsHeapPressureImpacting()) { + bytes += buffer->ComputeSize(size_fn); + } + } + } + return bytes; +} + +/*static*/ +bool HloLiveRange::BufferLivesOut(const HloBuffer& buffer, + const HloAliasAnalysis& alias_analysis, + const HloComputation* computation) { + if (alias_analysis.BufferLivesOut(buffer)) { + return true; + } + if (computation != nullptr) { + for (const HloValue* value : buffer.values()) { + for (const HloUse& use : value->GetUses()) { + if (use.instruction->parent() != computation) { + return true; + } + } + } + } + return false; +} + void HloLiveRange::NormalizeAliasedBuffers() { absl::flat_hash_map>> @@ -288,36 +403,26 @@ void HloLiveRange::CalculateBufferStartEndMap() { << definition_end_time; } - const InstructionValueSet& value_set_tree = - alias_analysis_.dataflow_analysis().GetInstructionValueSet( - &instruction); - - for (const auto& entry : value_set_tree) { - for (const HloValue* value : entry.second.values()) { - // The start time is only correct for the defining instruction. - if (value->defining_instruction() != &instruction) { - continue; - } - - auto [end_time, end_position] = - ComputeValueLiveRangeEnd(*value, definition_end_time); - LiveRangeBounds live_range{start_time, end_time, end_position}; - - // Readonly entry parameters (parameters that don't alias) live across - // whole computation. - const HloModule& module = *computation->parent(); - if (instruction.opcode() == HloOpcode::kParameter && - computation == module.entry_computation() && - !module.input_output_alias_config().ParameterHasAlias( - instruction.parameter_number(), value->index())) { - live_range.end = schedule_end_time(); - } else { - live_range.end = std::max(live_range.end, GetLastUsageTime(*value)); - } - - CHECK_LE(live_range.start, live_range.end) << instruction.ToString(); - CHECK(buffer_live_ranges_.insert({value, live_range}).second); + for (const HloValue* value : + GetValuesDefined(&instruction, alias_analysis_.dataflow_analysis())) { + auto [end_time, end_position] = + ComputeValueLiveRangeEnd(*value, definition_end_time); + LiveRangeBounds live_range{start_time, end_time, end_position}; + + // Readonly entry parameters (parameters that don't alias) live across + // whole computation. + const HloModule& module = *computation->parent(); + if (instruction.opcode() == HloOpcode::kParameter && + computation == module.entry_computation() && + !module.input_output_alias_config().ParameterHasAlias( + instruction.parameter_number(), value->index())) { + live_range.end = schedule_end_time(); + } else { + live_range.end = std::max(live_range.end, GetLastUsageTime(*value)); } + + CHECK_LE(live_range.start, live_range.end) << instruction.ToString(); + CHECK(buffer_live_ranges_.insert({value, live_range}).second); } } } diff --git a/third_party/xla/xla/hlo/utils/hlo_live_range.h b/third_party/xla/xla/hlo/utils/hlo_live_range.h index cc407ced4e42bd..92758777416b47 100644 --- a/third_party/xla/xla/hlo/utils/hlo_live_range.h +++ b/third_party/xla/xla/hlo/utils/hlo_live_range.h @@ -19,6 +19,7 @@ the License. #include #include #include +#include #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" @@ -26,9 +27,12 @@ the License. #include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "xla/hlo/analysis/hlo_alias_analysis.h" +#include "xla/hlo/analysis/hlo_dataflow_analysis.h" #include "xla/hlo/ir/dfs_hlo_visitor.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_schedule.h" +#include "xla/service/buffer_value.h" +#include "xla/service/hlo_buffer.h" #include "xla/service/hlo_value.h" namespace xla { @@ -46,6 +50,46 @@ class HloLiveRange { const HloComputation* computation, bool module_scoped_analysis = true, absl::flat_hash_set execution_threads = {}); + // Returns all HloValues defined by this instruction. + static std::vector GetValuesDefined( + const HloInstruction* instruction, const HloDataflowAnalysis& dataflow); + + // Returns the distinct physical HloBuffers newly allocated/defined by this + // instruction (excluding buffers forwarded or aliased from operands). + static std::vector GetBuffersDefined( + const HloInstruction* instruction, + const HloAliasAnalysis& alias_analysis); + + // Returns the total bytes defined by this instruction according to size_fn. + // Returns 0 for parameter instructions (parameters are attributed to + // computation start). + static int64_t GetBytesDefined(const HloInstruction* instruction, + const HloAliasAnalysis& alias_analysis, + const BufferValue::SizeFunction& size_fn); + + // Returns the distinct physical HloBuffers read by the operands of this + // instruction. + static std::vector GetBuffersUsed( + const HloInstruction* instruction, + const HloAliasAnalysis& alias_analysis); + + // Returns the total number of instruction reads across the computation for + // all values contained in this buffer. If computation is null, counts across + // all computations. + static int32_t GetTotalUsers(const HloBuffer& buffer, + const HloComputation* computation = nullptr); + + // Returns the total parameter bytes allocated at the start of the computation + // (matching HloLiveRange parameter attribution). + static int64_t GetParameterBytesAtStart( + const HloComputation& computation, const HloAliasAnalysis& alias_analysis, + const BufferValue::SizeFunction& size_fn); + + // Returns true if any value in this buffer lives out of the computation. + static bool BufferLivesOut(const HloBuffer& buffer, + const HloAliasAnalysis& alias_analysis, + const HloComputation* computation = nullptr); + // LogicalTime represents the time in a virtual clock. Each instruction has // one monotonically increasing logical time assigned according to the // schedule. diff --git a/third_party/xla/xla/hlo/utils/hlo_live_range_test.cc b/third_party/xla/xla/hlo/utils/hlo_live_range_test.cc index cc2a4a9963364e..e3bb28c339bff7 100644 --- a/third_party/xla/xla/hlo/utils/hlo_live_range_test.cc +++ b/third_party/xla/xla/hlo/utils/hlo_live_range_test.cc @@ -37,6 +37,7 @@ limitations under the License. #include "xla/hlo/parser/hlo_parser.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/literal_util.h" +#include "xla/service/buffer_value.h" #include "xla/service/hlo_value.h" #include "xla/shape.h" #include "xla/shape_util.h" @@ -1140,5 +1141,66 @@ ENTRY %entry { EXPECT_TRUE(hlo_live_range_->instruction_schedule().contains(neg0)); } +TEST_F(HloLiveRangeTest, HelpersTest) { + const char* hlo_string = R"( +HloModule module + +ENTRY entry { + param0 = f32[100] parameter(0) + param1 = f32[200] parameter(1) + const = f32[10] constant({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}) + add0 = f32[100] add(param0, param0) + add1 = f32[200] add(param1, param1) + ROOT root = (f32[100], f32[200]) tuple(add0, add1) +} +)"; + + ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr aa, + HloAliasAnalysis::Run(module_.get(), &alias_info_)); + + HloComputation* entry = module_->entry_computation(); + const HloInstruction* param0 = entry->parameter_instruction(0); + const HloInstruction* param1 = entry->parameter_instruction(1); + const HloInstruction* const_inst = entry->GetInstructionWithName("const"); + const HloInstruction* add0 = entry->GetInstructionWithName("add0"); + const HloInstruction* add1 = entry->GetInstructionWithName("add1"); + const HloInstruction* root = entry->root_instruction(); + + auto size_fn = [](const BufferValue& buffer) { + return ShapeUtil::ByteSizeOf(buffer.shape(), 8); + }; + + // 1. GetBuffersDefined + auto param0_buffers = HloLiveRange::GetBuffersDefined(param0, *aa); + EXPECT_EQ(param0_buffers.size(), 1); + auto param1_buffers = HloLiveRange::GetBuffersDefined(param1, *aa); + EXPECT_EQ(param1_buffers.size(), 1); + auto add0_buffers = HloLiveRange::GetBuffersDefined(add0, *aa); + EXPECT_EQ(add0_buffers.size(), 1); + + // 2. GetBytesDefined + EXPECT_EQ(HloLiveRange::GetBytesDefined(param0, *aa, size_fn), 0); + EXPECT_EQ(HloLiveRange::GetBytesDefined(add0, *aa, size_fn), 400); + EXPECT_EQ(HloLiveRange::GetBytesDefined(add1, *aa, size_fn), 800); + EXPECT_EQ(HloLiveRange::GetBytesDefined(const_inst, *aa, size_fn), 0); + + // 3. GetBuffersUsed + auto add0_used = HloLiveRange::GetBuffersUsed(add0, *aa); + EXPECT_EQ(add0_used.size(), 1); + EXPECT_EQ(add0_used[0], param0_buffers[0]); + + // 4. GetParameterBytesAtStart + EXPECT_EQ(HloLiveRange::GetParameterBytesAtStart(*entry, *aa, size_fn), 1200); + + // 5. GetTotalUsers + EXPECT_EQ(HloLiveRange::GetTotalUsers(*param0_buffers[0], entry), 2); + + // 6. BufferLivesOut + EXPECT_FALSE(HloLiveRange::BufferLivesOut(*param0_buffers[0], *aa, entry)); + auto root_buffers = HloLiveRange::GetBuffersDefined(root, *aa); + EXPECT_FALSE(root_buffers.empty()); + EXPECT_TRUE(HloLiveRange::BufferLivesOut(*root_buffers[0], *aa, entry)); +} } // namespace } // namespace xla diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index f1a7f7134f2de0..9e402b8148fed0 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -4059,7 +4059,9 @@ cc_library( ":buffer_value", ":hlo_value", "//xla:xla_data_proto_cc", + "//xla/hlo/ir:hlo", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", diff --git a/third_party/xla/xla/service/buffer_assignment.h b/third_party/xla/xla/service/buffer_assignment.h index bb898e2f1fb823..6d69deac524334 100644 --- a/third_party/xla/xla/service/buffer_assignment.h +++ b/third_party/xla/xla/service/buffer_assignment.h @@ -701,11 +701,7 @@ class BufferAssignment { int64_t HloBufferSize(const HloBuffer& buffer) { auto [it, inserted] = cached_buffer_sizes_.try_emplace(buffer.id()); if (inserted) { - int64_t result = 0; - for (const HloValue* value : buffer.values()) { - result = std::max(result, buffer_size_(*value)); - } - it->second = result; + it->second = buffer.ComputeSize(buffer_size_); } return it->second; } diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc index 999ffacb383873..847bcd49250dc4 100644 --- a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc +++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc @@ -347,7 +347,7 @@ absl::Status HeapSimulator::RunComputation( if (!buffer_live_ranges.contains(value)) { continue; } - if (IgnoreBuffer(value)) { + if (!IsHeapPressureImpacting(value)) { continue; } @@ -380,10 +380,7 @@ absl::Status HeapSimulator::RunComputation( // Populate buffer sizes with the maximum size of the constituent HloValues. for (const HloBuffer& buffer : alias_analysis.buffers()) { - int64_t size = 0; - for (const HloValue* value : buffer.values()) { - size = std::max(size, (*size_fn_)(*value)); - } + int64_t size = buffer.ComputeSize(*size_fn_); const HloValue* first_value = nullptr; for (const HloValue* value : buffer.values()) { buffer_groups_.emplace(value, size); @@ -443,7 +440,7 @@ absl::Status HeapSimulator::RunComputation( continue; } - if (IgnoreBuffer(operand_value)) { + if (!IsHeapPressureImpacting(operand_value)) { continue; } @@ -519,17 +516,9 @@ HeapSimulator::HeapSimulator( HeapSimulator::~HeapSimulator() {} -bool HeapSimulator::IgnoreBuffer(const HloValue* buffer) const { - // Buffers for constants are ignored unless the alloc_constants option is - // set. Also ignore buffers that we're not meant to assign. - // - // TODO(b/32248867): For consistency, constants should get allocations. - if (!options_.alloc_constants && - buffer->instruction()->opcode() == HloOpcode::kConstant) { - return true; - } - return options_.buffers_to_assign != nullptr && - !options_.buffers_to_assign->contains(buffer); +bool HeapSimulator::IsHeapPressureImpacting(const HloValue* buffer) const { + return HloBuffer::IsHeapPressureImpacting(*buffer, options_.alloc_constants, + options_.buffers_to_assign); } // Alloc always calls the underlying heap algorithm. diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.h b/third_party/xla/xla/service/heap_simulator/heap_simulator.h index beb04f68c64c2f..187e06dd731fb5 100644 --- a/third_party/xla/xla/service/heap_simulator/heap_simulator.h +++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.h @@ -214,7 +214,10 @@ class HeapSimulator { const HloAliasAnalysis& alias_analysis, const AliasInfo* alias_info, HloLiveRange* live_range); - bool IgnoreBuffer(const HloValue* buffer) const; + // Returns whether the buffer should be allocated space in the heap simulation + // (excludes constants unless alloc_constants is set, and respects the + // buffers_to_assign filter). + bool IsHeapPressureImpacting(const HloValue* buffer) const; void Alloc(const HloValue* buffer, const HloInstruction* instruction); void Free(const HloValue* buffer, const HloInstruction* instruction); // ShareBuffer indicates that a new buffer is defined and it has to be the diff --git a/third_party/xla/xla/service/hlo_buffer.cc b/third_party/xla/xla/service/hlo_buffer.cc index 7e2ba451f6a87f..7fe0f5041828d2 100644 --- a/third_party/xla/xla/service/hlo_buffer.cc +++ b/third_party/xla/xla/service/hlo_buffer.cc @@ -16,14 +16,19 @@ limitations under the License. #include "xla/service/hlo_buffer.h" #include +#include #include #include #include #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_set.h" #include "absl/log/check.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/service/buffer_value.h" #include "xla/service/hlo_value.h" namespace xla { @@ -73,4 +78,33 @@ std::ostream& operator<<(std::ostream& out, const HloBuffer& buffer) { return out; } +int64_t HloBuffer::ComputeSize(const BufferValue::SizeFunction& size_fn) const { + int64_t max_size = 0; + for (const HloValue* value : values_) { + max_size = std::max(max_size, size_fn(*value)); + } + return max_size; +} + +bool HloBuffer::IsHeapPressureImpacting( + const HloValue& value, bool alloc_constants, + const absl::flat_hash_set* buffers_to_assign) { + if (!alloc_constants && + value.instruction()->opcode() == HloOpcode::kConstant) { + return false; + } + if (buffers_to_assign != nullptr && !buffers_to_assign->contains(&value)) { + return false; + } + return true; +} + +bool HloBuffer::IsHeapPressureImpacting( + bool alloc_constants, + const absl::flat_hash_set* buffers_to_assign) const { + return absl::c_any_of(values_, [&](const HloValue* value) { + return IsHeapPressureImpacting(*value, alloc_constants, buffers_to_assign); + }); +} + } // namespace xla diff --git a/third_party/xla/xla/service/hlo_buffer.h b/third_party/xla/xla/service/hlo_buffer.h index 2ee0887ee76a08..4be60c9097d251 100644 --- a/third_party/xla/xla/service/hlo_buffer.h +++ b/third_party/xla/xla/service/hlo_buffer.h @@ -22,6 +22,7 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_set.h" #include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/statusor.h" @@ -97,6 +98,25 @@ class HloBuffer { // Return all values contained in this buffer. const std::vector& values() const { return values_; } + // Computes the physical size of the buffer as the maximum size of its + // constituent HloValues according to the given size function. + int64_t ComputeSize(const BufferValue::SizeFunction& size_fn) const; + + // Returns whether this value impacts dynamic heap allocation pressure (e.g. + // not an embedded constant or a buffer excluded by allocation filter). + static bool IsHeapPressureImpacting( + const HloValue& value, bool alloc_constants = false, + const absl::flat_hash_set* buffers_to_assign = nullptr); + + // Returns whether this buffer impacts dynamic heap allocation pressure. + // Matches HeapSimulator's allocation logic: a buffer impacts heap pressure if + // any of its constituent values impacts heap pressure (i.e. is not ignored + // by HeapSimulator). A buffer is only exempt if all of its constituent values + // are ignored (e.g. purely constants or outside buffers_to_assign). + bool IsHeapPressureImpacting(bool alloc_constants = false, + const absl::flat_hash_set* + buffers_to_assign = nullptr) const; + // Memory space color. Used to indicate the memory space that the hlo buffer // needs to live in. absl::StatusOr color() const { From 3da672680435d6bf026e65783d55c6e5b7dd6f4c Mon Sep 17 00:00:00 2001 From: Bill Varcho Date: Wed, 2 Sep 2026 19:12:33 -0700 Subject: [PATCH 13/15] [SDY][Bug-Fix][re-land] Fix Shardy shard_map export for nested ManualComputationOp across calls. PiperOrigin-RevId: 975451427 --- .../stablehlo_round_trip/shard_map_export.cc | 61 ++++++++++++------- ...stablehlo_round_trip_shard_map_export.mlir | 40 ++++++++++++ 2 files changed, 78 insertions(+), 23 deletions(-) diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc index 6a072cab73678b..71b6c94c5a3e40 100644 --- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc +++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc @@ -184,16 +184,24 @@ void setOpManualAxes(Operation* op, ManualAxesAttr manualAxes, op->setAttr(kManualAxes, manualAxes); } -void setFuncManualAxesRecursively(FuncOp funcOp, ManualAxesAttr manualAxes, - Attribute meshOrRef, - const mlir::SymbolTable& symbolTable); - -mlir::WalkResult setManualAxes(Operation* op, ManualAxesAttr manualAxes, - Attribute meshOrRef, - const mlir::SymbolTable& symbolTable) { - if (mlir::isa(op)) { - // Skip `ManualComputationOp`s and their nested operations, they will - // be handled separately. +void setManualAxesForOpsInBody( + ManualComputationOp op, const mlir::SymbolTable& symbolTable, + ManualComputationToParentManualAxes& parentManualCompAxes); + +void setFuncManualAxesRecursively( + FuncOp funcOp, ManualAxesAttr manualAxes, Attribute meshOrRef, + const mlir::SymbolTable& symbolTable, + ManualComputationToParentManualAxes& parentManualCompAxes); + +mlir::WalkResult setManualAxes( + Operation* op, ManualAxesAttr manualAxes, Attribute meshOrRef, + const mlir::SymbolTable& symbolTable, + ManualComputationToParentManualAxes& parentManualCompAxes) { + if (auto manualCompOp = mlir::dyn_cast(op)) { + // Record parent manual axes for this manualCompOp and process its body. + parentManualCompAxes[manualCompOp].assign(manualAxes.getValue().begin(), + manualAxes.getValue().end()); + setManualAxesForOpsInBody(manualCompOp, symbolTable, parentManualCompAxes); return mlir::WalkResult::skip(); } if (!mlir::isa(op)) { @@ -202,14 +210,16 @@ mlir::WalkResult setManualAxes(Operation* op, ManualAxesAttr manualAxes, if (CallOp callOp = mlir::dyn_cast(op)) { FuncOp funcOp = symbolTable.lookup(callOp.getCallee()); CHECK(funcOp) << "Failed to lookup function: " << callOp.getCallee().str(); - setFuncManualAxesRecursively(funcOp, manualAxes, meshOrRef, symbolTable); + setFuncManualAxesRecursively(funcOp, manualAxes, meshOrRef, symbolTable, + parentManualCompAxes); } return mlir::WalkResult::advance(); } -void setFuncManualAxesRecursively(FuncOp funcOp, ManualAxesAttr manualAxes, - Attribute meshOrRef, - const mlir::SymbolTable& symbolTable) { +void setFuncManualAxesRecursively( + FuncOp funcOp, ManualAxesAttr manualAxes, Attribute meshOrRef, + const mlir::SymbolTable& symbolTable, + ManualComputationToParentManualAxes& parentManualCompAxes) { llvm::SmallVector funcArgAttrs; funcArgAttrs.reserve(funcOp.getNumArguments()); for (int argNum = 0; argNum < funcOp.getNumArguments(); argNum++) { @@ -250,15 +260,15 @@ void setFuncManualAxesRecursively(FuncOp funcOp, ManualAxesAttr manualAxes, // Walk in preorder of blocks in order to stop walks on manual computations. funcOp->walk([&](Operation* op) { - return setManualAxes(op, manualAxes, meshOrRef, symbolTable); + return setManualAxes(op, manualAxes, meshOrRef, symbolTable, + parentManualCompAxes); }); } // Sets the manual axes of all operations in `op`'s body. void setManualAxesForOpsInBody( - ManualComputationOp op, - const ManualComputationToParentManualAxes& parentManualCompAxes, - const mlir::SymbolTable& symbolTable) { + ManualComputationOp op, const mlir::SymbolTable& symbolTable, + ManualComputationToParentManualAxes& parentManualCompAxes) { TensorShardingAttr sharding = getFirstSharding(op); if (!sharding) { // If there are no in/out shardings, op.getManualAxes() must be empty. We do @@ -280,7 +290,8 @@ void setManualAxesForOpsInBody( // Set the manual axes of all operations in the body. op.getBody().front().walk( [&](Operation* opInBody) { - return setManualAxes(opInBody, manualAxesAttr, meshOrRef, symbolTable); + return setManualAxes(opInBody, manualAxesAttr, meshOrRef, symbolTable, + parentManualCompAxes); }); } @@ -483,12 +494,16 @@ class ShardMapExportPass // walk. module->walk([&](ManualComputationOp op) { if (auto parentOp = op->getParentOfType()) { - SmallVector& parentAxes = parentManualCompAxes[op]; - parentAxes = parentManualCompAxes[parentOp]; - parentAxes.insert(parentAxes.end(), parentOp.getManualAxes().begin(), + SmallVector parentAxes; + if (auto it = parentManualCompAxes.find(parentOp); + it != parentManualCompAxes.end()) { + parentAxes = it->second; + } + parentAxes.append(parentOp.getManualAxes().begin(), parentOp.getManualAxes().end()); + parentManualCompAxes[op] = std::move(parentAxes); } - setManualAxesForOpsInBody(op, parentManualCompAxes, symbolTable); + setManualAxesForOpsInBody(op, symbolTable, parentManualCompAxes); }); // Need to do a separate post order walk to inline the diff --git a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_round_trip_shard_map_export.mlir b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_round_trip_shard_map_export.mlir index 8a7d5aa3cbd6b4..5ca34525922de5 100644 --- a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_round_trip_shard_map_export.mlir +++ b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_round_trip_shard_map_export.mlir @@ -575,3 +575,43 @@ func.func private @bar(%arg0: tensor<8xi32>) -> tensor<8xi32> { // CHECK-NEXT: %1 = call @bar(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8xi32>) -> tensor<8xi32> // CHECK-NEXT: return %arg0 : tensor<4xi32> // CHECK-NEXT: } + +// ----- +sdy.mesh @mesh = <["a"=2, "b"=4]> + +// Tests that when a function containing a `ManualComputationOp` (callee) is +// called from within another `ManualComputationOp` (caller), the callee's +// `ManualComputationOp` correctly inherits parent manual axes from the caller. +// +// 1. Caller passes parent manual axes {"a"} to callee's inputs/outputs: +// CHECK-LABEL: func private @called_func_with_manual_comp( +// CHECK-SAME: %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>, xla.sdy.manual_axes = #sdy}) +// CHECK-SAME: -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>, xla.sdy.manual_axes = #sdy}) { +// CHECK-NEXT: %[[COPY:.*]] = mhlo.copy %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"b"}]>]>, xla.sdy.manual_axes = #sdy} : tensor<8x8xf32> +// +// 2. Full-to-shard inside callee inherits caller axis "a" + callee axis "b" -> manual axes {"a", "b"}: +// CHECK-NEXT: %[[FULL_TO_SHARD:.*]] = stablehlo.custom_call @SPMDFullToShardShape(%[[COPY]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x8xf32>) -> tensor<8x2xf32> +// +// 3. Inner body call runs with both axes manual: +// CHECK-NEXT: %[[CALL:.*]] = call @xla.sdy.inlinable_manual_computation_body(%[[FULL_TO_SHARD]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x2xf32>) -> tensor<8x2xf32> +// CHECK-NEXT: %[[COPY_1:.*]] = mhlo.copy %[[CALL]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : tensor<8x2xf32> +// +// 4. Shard-to-full returns to caller scope, with manual axis {"a"} preserved: +// CHECK-NEXT: %[[SHARD_TO_FULL:.*]] = stablehlo.custom_call @SPMDShardToFullShape(%[[COPY_1]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"b"}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x2xf32>) -> tensor<8x8xf32> +// CHECK-NEXT: return %[[SHARD_TO_FULL]] : tensor<8x8xf32> +func.func private @called_func_with_manual_comp(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> { + %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{}, {"b"}]>] out_shardings=[<@mesh, [{}, {"b"}]>] manual_axes={"b"} (%arg1: tensor<8x2xf32>) { + %1 = stablehlo.add %arg1, %arg1 : tensor<8x2xf32> + sdy.return %1 : tensor<8x2xf32> + } : (tensor<8x8xf32>) -> tensor<8x8xf32> + return %0 : tensor<8x8xf32> +} + +// CHECK-LABEL: func @manual_comp_calls_func_with_manual_comp( +func.func @manual_comp_calls_func_with_manual_comp(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> { + %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{"a"}, {}]>] out_shardings=[<@mesh, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<8x8xf32>) { + %1 = func.call @called_func_with_manual_comp(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32> + sdy.return %1 : tensor<8x8xf32> + } : (tensor<16x8xf32>) -> tensor<16x8xf32> + return %0 : tensor<16x8xf32> +} From b5a035efd671d3534cd3e4a0be5ffac705502add Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Longeri?= Date: Wed, 2 Sep 2026 20:02:05 -0700 Subject: [PATCH 14/15] [Mosaic] Expose global variables for special semaphores PiperOrigin-RevId: 975465130 --- third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h b/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h index a87e9cd7426b5a..ea72eb35a3be98 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h @@ -53,6 +53,11 @@ class TPUDialect; namespace mlir { namespace tpu { +inline constexpr StringRef kTpuVectorSubcoreGlobalBarrierName = + "__tpu_vector_subcore_global_barrier"; +inline constexpr StringRef kTpuScalarSubcoreGlobalBarrierName = + "__tpu_scalar_subcore_global_barrier"; + DEFINE_ABSL_STRINGIFY_FOR_ENUMS(); struct TpuTilingFlags { From 0f794b9658a6acd19c2ecdc61ec84d1fde367c16 Mon Sep 17 00:00:00 2001 From: Srishti Srivastava Date: Wed, 2 Sep 2026 22:23:27 -0700 Subject: [PATCH 15/15] Integrate StableHLO at openxla/stablehlo@63933593 PiperOrigin-RevId: 975511528 --- .../xla/third_party/stablehlo/temporary.patch | 859 ------------------ .../xla/third_party/stablehlo/workspace.bzl | 4 +- .../translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc | 24 +- .../xla/xla/mlir_hlo/mhlo/IR/hlo_ops.cc | 4 +- third_party/xla/xla/python/version.h | 3 +- 5 files changed, 25 insertions(+), 869 deletions(-) diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch index ebb9dda1528d7c..b1e2e8b7ff4ea0 100644 --- a/third_party/xla/third_party/stablehlo/temporary.patch +++ b/third_party/xla/third_party/stablehlo/temporary.patch @@ -15,863 +15,4 @@ # # This file is automatically generated by generate_patch tool. # Do not edit directly. -diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md ---- stablehlo/docs/spec.md -+++ stablehlo/docs/spec.md -@@ -2493,7 +2493,9 @@ - * `num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1`. - * (C26) `rank(result) = N`. - * If the operation uses non-quantized tensors: -- * (C27) `element_type(lhs) = element_type(rhs) = element_type(result)`. -+ * (C27) `element_type(lhs) = element_type(rhs) = element_type(result)` or -+ (`is_fp8(element_type(lhs))` and `is_fp8(element_type(rhs))` and -+ `element_type(result) = element_type(lhs)`). - * If the operation uses quantized tensors: - * (C28) `is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)`. - * (C29) If `is_per_axis_quantized(rhs)`, -@@ -3163,7 +3165,9 @@ - * `num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1`. - * (C26) `rank(result) = N`. - * If the operation uses non-quantized tensors: -- * (C27) `element_type(lhs) = element_type(rhs) = element_type(result)`. -+ * (C27) `element_type(lhs) = element_type(rhs) = element_type(result)` or -+ (`is_fp8(element_type(lhs))` and `is_fp8(element_type(rhs))` and -+ `element_type(result) = element_type(lhs)`). - * If the operation uses quantized tensors: - * (C28) `is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)`. - * (C29) If `is_per_axis_quantized(rhs)`, -@@ -7654,6 +7658,12 @@ - * `is_quantized(x: Value | Placeholder | Type) -> Value` is a shortcut for - `is_quantized_tensor_element_type(x)`. - -+* `is_fp8(x: Value | Placeholder | Type) -> Value` returns `true` if `x` is one -+of `Float8E4M3Type`, `Float8E4M3FNType`, `Float8E4M3B11FNUZType`, -+`Float8E4M3FNUZType`, `Float8E5M2Type`, `Float8E5M2FNUZType`, `Float8E3M4Type`, -+or `Float8E8M0FNUType`. If `x` is a value or placeholder, this function is a -+shortcut for `is_fp8(type(x))`. -+ - * `is_type_name(x: Value | Placeholder | Type) -> Value`. Available for all - types. For example, `is_float(x)` returns `true` if `x` is a `FloatType`. - If `x` is a value or placeholder, this function is a shortcut for -diff --ruN a/stablehlo/stablehlo/dialect/Serialization.cpp b/stablehlo/stablehlo/dialect/Serialization.cpp ---- stablehlo/stablehlo/dialect/Serialization.cpp -+++ stablehlo/stablehlo/dialect/Serialization.cpp -@@ -24,6 +24,7 @@ - #include "mlir/IR/Location.h" - #include "mlir/IR/MLIRContext.h" - #include "mlir/IR/OwningOpRef.h" -+#include "mlir/IR/Verifier.h" - #include "mlir/Parser/Parser.h" - #include "mlir/Pass/PassManager.h" - #include "mlir/Support/LLVM.h" -@@ -43,11 +44,18 @@ - bool allowOtherDialects) { - MLIRContext* context = module.getContext(); - -+ // Only enable verifier in debug builds. -+ bool enableVerifier = false; -+#ifndef NDEBUG -+ enableVerifier = true; -+#endif -+ - // Convert StableHLO --> VHLO. - // If allowOtherDialects is true, we will allow other dialects to be present - // in the module, otherwise will fail if there are any other dialects present. - { - PassManager pm(context); -+ pm.enableVerifier(enableVerifier); - StablehloLegalizeToVhloPassOptions options; - options.allowOtherDialects = allowOtherDialects; - pm.addPass(stablehlo::createStablehloLegalizeToVhloPass(options)); -@@ -61,6 +69,7 @@ - // target version failures. - { - PassManager pm(context); -+ pm.enableVerifier(enableVerifier); - pm.addPass(stablehlo::createVhloToVersionPass({targetVersion.str()})); - if (!succeeded(pm.run(module))) { - return failure(); -@@ -81,7 +90,15 @@ - OwningOpRef deserializePortableArtifact(StringRef sourceStr, - MLIRContext* context) { - context->loadDialect(); -- auto module = parseSourceString(sourceStr, context); -+ -+ // Only enable verifier in debug builds. -+ bool enableVerifier = false; -+#ifndef NDEBUG -+ enableVerifier = true; -+#endif -+ -+ ParserConfig config(context, /*verifyAfterParse=*/enableVerifier); -+ auto module = parseSourceString(sourceStr, config); - if (!module) { - emitError(UnknownLoc::get(context)) - << "failed to deserialize portable artifact using StableHLO_v" -@@ -91,8 +108,13 @@ - - // Convert VHLO --> VHLO(current) --> StableHLO - PassManager pm(context); -+ pm.enableVerifier(enableVerifier); - createStablehloDeserializePipeline(pm); - if (!succeeded(pm.run(*module))) { -+ return nullptr; -+ } -+ -+ if (failed(verify(*module))) { - return nullptr; - } - -diff --ruN a/stablehlo/stablehlo/dialect/TypeInference.cpp b/stablehlo/stablehlo/dialect/TypeInference.cpp ---- stablehlo/stablehlo/dialect/TypeInference.cpp -+++ stablehlo/stablehlo/dialect/TypeInference.cpp -@@ -86,6 +86,13 @@ - // Utils for quantization specific verifications - //===----------------------------------------------------------------------===// - -+bool isFp8Type(mlir::Type type) { -+ return llvm::isa(type); -+} -+ - template - bool allQuantized(ArrayRef typeRange) { - return llvm::all_of( -@@ -2221,11 +2228,14 @@ - // convolution_c27 - if (!anyQuantized({rankedLhsType, rankedRhsType}) && - !isCompatibleForHloTypeInference(rankedLhsType.getElementType(), -- rankedRhsType.getElementType())) -+ rankedRhsType.getElementType()) && -+ !(isFp8Type(rankedLhsType.getElementType()) && -+ isFp8Type(rankedRhsType.getElementType()))) { - return emitOptionalError( - location, "expects lhs and rhs to have compatible element type. Got: ", - rankedLhsType.getElementType(), " and ", - rankedRhsType.getElementType()); -+ } - - if (failed(verifyConvolutionAttributes( - location, lhsType, rhsType, inputBatchDimension, -@@ -2548,7 +2558,9 @@ - // dynamic_conv_c27 - if (!anyQuantized({rankedLhsType, rankedRhsType}) && - !isCompatibleForHloTypeInference(rankedLhsType.getElementType(), -- rankedRhsType.getElementType())) -+ rankedRhsType.getElementType()) && -+ !(isFp8Type(rankedLhsType.getElementType()) && -+ isFp8Type(rankedRhsType.getElementType()))) - return emitOptionalError( - location, "expects lhs and rhs to have compatible element type. Got: ", - rankedLhsType.getElementType(), " and ", -diff --ruN a/stablehlo/stablehlo/dialect/Version.h b/stablehlo/stablehlo/dialect/Version.h ---- stablehlo/stablehlo/dialect/Version.h -+++ stablehlo/stablehlo/dialect/Version.h -@@ -38,7 +38,7 @@ - static FailureOr fromString(llvm::StringRef versionRef); - - /// Return a Version representing the current VHLO dialect version. -- static Version getCurrentVersion() { return Version(1, 19, 0); } -+ static Version getCurrentVersion() { return Version(1, 20, 0); } - - /// Return a Version representing the minimum supported VHLO dialect version. - static Version getMinimumVersion() { return Version(0, 9, 0); } -diff --ruN a/stablehlo/stablehlo/dialect/VhloDialect.td b/stablehlo/stablehlo/dialect/VhloDialect.td ---- stablehlo/stablehlo/dialect/VhloDialect.td -+++ stablehlo/stablehlo/dialect/VhloDialect.td -@@ -58,6 +58,7 @@ - 1.17.0: Add `future` type support to `custom_call` op. - 1.18.0: Add `result_tilings` attribute to `custom_call` op. - 1.19.0: Add CollectiveReduceOp. -+ 1.20.0: Allow mixed fp8 operands in `convolution` and `dynamic_conv` ops. - }]; - - let useDefaultAttributePrinterParser = 0; -diff --ruN a/stablehlo/stablehlo/dialect/VhloOps.cpp b/stablehlo/stablehlo/dialect/VhloOps.cpp ---- stablehlo/stablehlo/dialect/VhloOps.cpp -+++ stablehlo/stablehlo/dialect/VhloOps.cpp -@@ -359,11 +359,44 @@ - return success(); - } - -+bool isVhloFp8Type(Type type) { -+ return isa(type); -+} -+ -+LogicalResult verifyConstraint_1_20_0(mlir::Operation* op, -+ Version targetVersion) { -+ if (targetVersion < Version(1, 20, 0)) { -+ if (op->getNumOperands() < 2) { -+ return failure(); -+ } -+ Type lhsElementType = getVhloElementType(op->getOperand(0).getType()); -+ Type rhsElementType = getVhloElementType(op->getOperand(1).getType()); -+ if (lhsElementType != rhsElementType && isVhloFp8Type(lhsElementType) && -+ isVhloFp8Type(rhsElementType)) { -+ return failure(); -+ } -+ } -+ return success(); -+} -+ - } // namespace - - LogicalResult AllReduceOpV1::validateConstraint(mlir::Operation* op, - Version targetVersion) { - return verifyConstraint_0_17_0(op, targetVersion); -+} -+ -+LogicalResult ConvolutionOpV1::validateConstraint(mlir::Operation* op, -+ Version targetVersion) { -+ return verifyConstraint_1_20_0(op, targetVersion); -+} -+ -+LogicalResult DynamicConvOpV2::validateConstraint(mlir::Operation* op, -+ Version targetVersion) { -+ return verifyConstraint_1_20_0(op, targetVersion); - } - - LogicalResult ReduceOpV1::validateConstraint(mlir::Operation* op, -diff --ruN a/stablehlo/stablehlo/dialect/VhloOps.td b/stablehlo/stablehlo/dialect/VhloOps.td ---- stablehlo/stablehlo/dialect/VhloOps.td -+++ stablehlo/stablehlo/dialect/VhloOps.td -@@ -390,7 +390,8 @@ - let results = (outs VHLO_AnyType:$result); - } - --def VHLO_ConvolutionOpV1 : VHLO_Op<"convolution_v1", "0.9.0", "current"> { -+def VHLO_ConvolutionOpV1 : VHLO_Op<"convolution_v1", "0.9.0", "current", -+ [DeclareOpInterfaceMethods]> { - let arguments = (ins - VHLO_AnyType:$lhs, - VHLO_AnyType:$rhs, -@@ -572,7 +573,8 @@ - - // Padding should be specified as an operand only, not an attribute. - // Remove `d_padding` and convert `padding` to an operand. --def VHLO_DynamicConvOpV2 : VHLO_Op<"dynamic_conv_v2", "0.20.0", "current"> { -+def VHLO_DynamicConvOpV2 : VHLO_Op<"dynamic_conv_v2", "0.20.0", "current", -+ [DeclareOpInterfaceMethods]> { - let arguments = (ins - VHLO_AnyType:$lhs, - VHLO_AnyType:$rhs, -diff --ruN a/stablehlo/stablehlo/tests/TestUtils.cpp b/stablehlo/stablehlo/tests/TestUtils.cpp ---- stablehlo/stablehlo/tests/TestUtils.cpp -+++ stablehlo/stablehlo/tests/TestUtils.cpp -@@ -68,6 +68,32 @@ - } - }; - -+struct BroadcastIfNeededPattern : public RewritePattern { -+ explicit BroadcastIfNeededPattern(MLIRContext* context) -+ : RewritePattern("hlo_test_broadcast.broadcast_if_needed", 1, context) {} -+ LogicalResult matchAndRewrite(Operation* op, -+ PatternRewriter& rewriter) const override { -+ if (op->getNumOperands() < 1) return failure(); -+ Value input = op->getOperand(0); -+ -+ SmallVector broadcastDimensions; -+ if (auto bcastDimsAttr = -+ op->getAttrOfType("broadcast_dimensions")) { -+ broadcastDimensions = llvm::to_vector(bcastDimsAttr.asArrayRef()); -+ } -+ -+ auto targetShape = stablehlo::getDimensions(op->getResult(0)); -+ if (failed(targetShape)) return failure(); -+ -+ auto broadcastedVal = stablehlo::broadcastIfNeeded( -+ rewriter, input, *targetShape, broadcastDimensions); -+ if (failed(broadcastedVal)) return failure(); -+ -+ rewriter.replaceOp(op, *broadcastedVal); -+ return success(); -+ } -+}; -+ - struct InferReturnTypesPattern : public RewritePattern { - explicit InferReturnTypesPattern(MLIRContext* context) - : RewritePattern("hlo_test_infer.get_return_types", 1, context) {} -@@ -222,6 +248,7 @@ - LogicalResult initialize(MLIRContext* context) override { - RewritePatternSet patterns(context); - patterns.add(context); -+ patterns.add(context); - patterns_ = std::move(patterns); - return success(); - } -diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir ---- stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir -+++ stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir -@@ -2141,13 +2141,16 @@ - // CHECK: %[[DIVIDE_19:.*]] = stablehlo.divide %[[MULTIPLY_9]], %[[SINE_0]] : tensor - // CHECK: %[[SUBTRACT_10:.*]] = stablehlo.subtract %[[SUBTRACT_9]], %[[DIVIDE_19]] : tensor - // CHECK: %[[SELECT_1:.*]] = stablehlo.select %[[COMPARE_0]], %[[SUBTRACT_10]], %[[SUBTRACT_9]] : tensor, tensor --// CHECK: %[[COMPARE_1:.*]] = stablehlo.compare LE, %[[ARG0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_1:.*]] = stablehlo.compare EQ, %[[ARG0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_2:.*]] = stablehlo.compare LT, %[[ARG0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor - // CHECK: %[[FLOOR_1:.*]] = stablehlo.floor %[[ARG0]] : tensor --// CHECK: %[[COMPARE_2:.*]] = stablehlo.compare EQ, %[[ARG0]], %[[FLOOR_1]] : (tensor, tensor) -> tensor --// CHECK: %[[AND_0:.*]] = stablehlo.and %[[COMPARE_1]], %[[COMPARE_2]] : tensor -+// CHECK: %[[COMPARE_3:.*]] = stablehlo.compare EQ, %[[ARG0]], %[[FLOOR_1]] : (tensor, tensor) -> tensor -+// CHECK: %[[AND_0:.*]] = stablehlo.and %[[COMPARE_2]], %[[COMPARE_3]] : tensor - // CHECK: %[[CONSTANT_25:.*]] = stablehlo.constant dense<0x7FF8000000000000> : tensor - // CHECK: %[[SELECT_2:.*]] = stablehlo.select %[[AND_0]], %[[CONSTANT_25]], %[[SELECT_1]] : tensor, tensor --// CHECK: return %[[SELECT_2]] : tensor -+// CHECK: %[[CONSTANT_26:.*]] = stablehlo.constant dense<0xFFF0000000000000> : tensor -+// CHECK: %[[SELECT_3:.*]] = stablehlo.select %[[COMPARE_1]], %[[CONSTANT_26]], %[[SELECT_2]] : tensor, tensor -+// CHECK: return %[[SELECT_3]] : tensor - // CHECK: } - func.func @digamma_f64(%arg : tensor) -> tensor { - %1 = chlo.digamma %arg : tensor -> tensor -@@ -2254,13 +2257,16 @@ - // CHECK: %[[DIVIDE_19:.*]] = stablehlo.divide %[[MULTIPLY_9]], %[[SINE_0]] : tensor - // CHECK: %[[SUBTRACT_10:.*]] = stablehlo.subtract %[[SUBTRACT_9]], %[[DIVIDE_19]] : tensor - // CHECK: %[[SELECT_1:.*]] = stablehlo.select %[[COMPARE_0]], %[[SUBTRACT_10]], %[[SUBTRACT_9]] : tensor, tensor --// CHECK: %[[COMPARE_1:.*]] = stablehlo.compare LE, %[[ARG0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_1:.*]] = stablehlo.compare EQ, %[[ARG0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_2:.*]] = stablehlo.compare LT, %[[ARG0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor - // CHECK: %[[FLOOR_1:.*]] = stablehlo.floor %[[ARG0]] : tensor --// CHECK: %[[COMPARE_2:.*]] = stablehlo.compare EQ, %[[ARG0]], %[[FLOOR_1]] : (tensor, tensor) -> tensor --// CHECK: %[[AND_0:.*]] = stablehlo.and %[[COMPARE_1]], %[[COMPARE_2]] : tensor -+// CHECK: %[[COMPARE_3:.*]] = stablehlo.compare EQ, %[[ARG0]], %[[FLOOR_1]] : (tensor, tensor) -> tensor -+// CHECK: %[[AND_0:.*]] = stablehlo.and %[[COMPARE_2]], %[[COMPARE_3]] : tensor - // CHECK: %[[CONSTANT_25:.*]] = stablehlo.constant dense<0x7FC00000> : tensor - // CHECK: %[[SELECT_2:.*]] = stablehlo.select %[[AND_0]], %[[CONSTANT_25]], %[[SELECT_1]] : tensor, tensor --// CHECK: return %[[SELECT_2]] : tensor -+// CHECK: %[[CONSTANT_26:.*]] = stablehlo.constant dense<0xFF800000> : tensor -+// CHECK: %[[SELECT_3:.*]] = stablehlo.select %[[COMPARE_1]], %[[CONSTANT_26]], %[[SELECT_2]] : tensor, tensor -+// CHECK: return %[[SELECT_3]] : tensor - // CHECK: } - func.func @digamma_f32(%arg : tensor) -> tensor { - %1 = chlo.digamma %arg : tensor -> tensor -@@ -2368,13 +2374,16 @@ - // CHECK: %[[DIVIDE_19:.*]] = stablehlo.divide %[[MULTIPLY_9]], %[[SINE_0]] : tensor - // CHECK: %[[SUBTRACT_10:.*]] = stablehlo.subtract %[[SUBTRACT_9]], %[[DIVIDE_19]] : tensor - // CHECK: %[[SELECT_1:.*]] = stablehlo.select %[[COMPARE_0]], %[[SUBTRACT_10]], %[[SUBTRACT_9]] : tensor, tensor --// CHECK: %[[COMPARE_1:.*]] = stablehlo.compare LE, %[[CONVERT_0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_1:.*]] = stablehlo.compare EQ, %[[CONVERT_0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_2:.*]] = stablehlo.compare LT, %[[CONVERT_0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor - // CHECK: %[[FLOOR_1:.*]] = stablehlo.floor %[[CONVERT_0]] : tensor --// CHECK: %[[COMPARE_2:.*]] = stablehlo.compare EQ, %[[CONVERT_0]], %[[FLOOR_1]] : (tensor, tensor) -> tensor --// CHECK: %[[AND_0:.*]] = stablehlo.and %[[COMPARE_1]], %[[COMPARE_2]] : tensor -+// CHECK: %[[COMPARE_3:.*]] = stablehlo.compare EQ, %[[CONVERT_0]], %[[FLOOR_1]] : (tensor, tensor) -> tensor -+// CHECK: %[[AND_0:.*]] = stablehlo.and %[[COMPARE_2]], %[[COMPARE_3]] : tensor - // CHECK: %[[CONSTANT_25:.*]] = stablehlo.constant dense<0x7FC00000> : tensor - // CHECK: %[[SELECT_2:.*]] = stablehlo.select %[[AND_0]], %[[CONSTANT_25]], %[[SELECT_1]] : tensor, tensor --// CHECK: %[[CONVERT_1:.*]] = stablehlo.convert %[[SELECT_2]] : (tensor) -> tensor -+// CHECK: %[[CONSTANT_26:.*]] = stablehlo.constant dense<0xFF800000> : tensor -+// CHECK: %[[SELECT_3:.*]] = stablehlo.select %[[COMPARE_1]], %[[CONSTANT_26]], %[[SELECT_2]] : tensor, tensor -+// CHECK: %[[CONVERT_1:.*]] = stablehlo.convert %[[SELECT_3]] : (tensor) -> tensor - // CHECK: return %[[CONVERT_1]] : tensor - // CHECK: } - func.func @digamma_f16(%arg : tensor) -> tensor { -@@ -2942,17 +2951,20 @@ - // CHECK: %[[DIVIDE_32:.*]] = stablehlo.divide %[[MULTIPLY_52]], %[[SINE_1]] : tensor - // CHECK: %[[SUBTRACT_18:.*]] = stablehlo.subtract %[[SUBTRACT_17]], %[[DIVIDE_32]] : tensor - // CHECK: %[[SELECT_12:.*]] = stablehlo.select %[[COMPARE_12]], %[[SUBTRACT_18]], %[[SUBTRACT_17]] : tensor, tensor --// CHECK: %[[COMPARE_13:.*]] = stablehlo.compare LE, %[[ARG1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_13:.*]] = stablehlo.compare EQ, %[[ARG1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_14:.*]] = stablehlo.compare LT, %[[ARG1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor - // CHECK: %[[FLOOR_5:.*]] = stablehlo.floor %[[ARG1]] : tensor --// CHECK: %[[COMPARE_14:.*]] = stablehlo.compare EQ, %[[ARG1]], %[[FLOOR_5]] : (tensor, tensor) -> tensor --// CHECK: %[[AND_3:.*]] = stablehlo.and %[[COMPARE_13]], %[[COMPARE_14]] : tensor -+// CHECK: %[[COMPARE_15:.*]] = stablehlo.compare EQ, %[[ARG1]], %[[FLOOR_5]] : (tensor, tensor) -> tensor -+// CHECK: %[[AND_3:.*]] = stablehlo.and %[[COMPARE_14]], %[[COMPARE_15]] : tensor - // CHECK: %[[CONSTANT_96:.*]] = stablehlo.constant dense<0x7FC00000> : tensor - // CHECK: %[[SELECT_13:.*]] = stablehlo.select %[[AND_3]], %[[CONSTANT_96]], %[[SELECT_12]] : tensor, tensor --// CHECK: %[[SELECT_14:.*]] = stablehlo.select %[[COMPARE_11]], %[[SELECT_13]], %[[MULTIPLY_42]] : tensor, tensor -+// CHECK: %[[CONSTANT_98:.*]] = stablehlo.constant dense<0xFF800000> : tensor -+// CHECK: %[[SELECT_16:.*]] = stablehlo.select %[[COMPARE_13]], %[[CONSTANT_98]], %[[SELECT_13]] : tensor, tensor -+// CHECK: %[[SELECT_14:.*]] = stablehlo.select %[[COMPARE_11]], %[[SELECT_16]], %[[MULTIPLY_42]] : tensor, tensor - // CHECK: %[[FLOOR_6:.*]] = stablehlo.floor %[[ARG0]] : tensor --// CHECK: %[[COMPARE_15:.*]] = stablehlo.compare NE, %[[ARG0]], %[[FLOOR_6]] : (tensor, tensor) -> tensor --// CHECK: %[[COMPARE_16:.*]] = stablehlo.compare LT, %[[ARG0]], %[[CONSTANT_70]] : (tensor, tensor) -> tensor --// CHECK: %[[OR_0:.*]] = stablehlo.or %[[COMPARE_15]], %[[COMPARE_16]] : tensor -+// CHECK: %[[COMPARE_16:.*]] = stablehlo.compare NE, %[[ARG0]], %[[FLOOR_6]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_17:.*]] = stablehlo.compare LT, %[[ARG0]], %[[CONSTANT_70]] : (tensor, tensor) -> tensor -+// CHECK: %[[OR_0:.*]] = stablehlo.or %[[COMPARE_16]], %[[COMPARE_17]] : tensor - // CHECK: %[[CONSTANT_97:.*]] = stablehlo.constant dense<0x7FC00000> : tensor - // CHECK: %[[SELECT_15:.*]] = stablehlo.select %[[OR_0]], %[[CONSTANT_97]], %[[SELECT_14]] : tensor, tensor - // CHECK: return %[[SELECT_15]] : tensor -@@ -3332,17 +3344,20 @@ - // CHECK: %[[DIVIDE_32:.*]] = stablehlo.divide %[[MULTIPLY_52]], %[[SINE_1]] : tensor - // CHECK: %[[SUBTRACT_18:.*]] = stablehlo.subtract %[[SUBTRACT_17]], %[[DIVIDE_32]] : tensor - // CHECK: %[[SELECT_12:.*]] = stablehlo.select %[[COMPARE_12]], %[[SUBTRACT_18]], %[[SUBTRACT_17]] : tensor, tensor --// CHECK: %[[COMPARE_13:.*]] = stablehlo.compare LE, %[[ARG1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_13:.*]] = stablehlo.compare EQ, %[[ARG1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_14:.*]] = stablehlo.compare LT, %[[ARG1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor - // CHECK: %[[FLOOR_5:.*]] = stablehlo.floor %[[ARG1]] : tensor --// CHECK: %[[COMPARE_14:.*]] = stablehlo.compare EQ, %[[ARG1]], %[[FLOOR_5]] : (tensor, tensor) -> tensor --// CHECK: %[[AND_3:.*]] = stablehlo.and %[[COMPARE_13]], %[[COMPARE_14]] : tensor -+// CHECK: %[[COMPARE_15:.*]] = stablehlo.compare EQ, %[[ARG1]], %[[FLOOR_5]] : (tensor, tensor) -> tensor -+// CHECK: %[[AND_3:.*]] = stablehlo.and %[[COMPARE_14]], %[[COMPARE_15]] : tensor - // CHECK: %[[CONSTANT_96:.*]] = stablehlo.constant dense<0x7FF8000000000000> : tensor - // CHECK: %[[SELECT_13:.*]] = stablehlo.select %[[AND_3]], %[[CONSTANT_96]], %[[SELECT_12]] : tensor, tensor --// CHECK: %[[SELECT_14:.*]] = stablehlo.select %[[COMPARE_11]], %[[SELECT_13]], %[[MULTIPLY_42]] : tensor, tensor -+// CHECK: %[[CONSTANT_98:.*]] = stablehlo.constant dense<0xFFF0000000000000> : tensor -+// CHECK: %[[SELECT_16:.*]] = stablehlo.select %[[COMPARE_13]], %[[CONSTANT_98]], %[[SELECT_13]] : tensor, tensor -+// CHECK: %[[SELECT_14:.*]] = stablehlo.select %[[COMPARE_11]], %[[SELECT_16]], %[[MULTIPLY_42]] : tensor, tensor - // CHECK: %[[FLOOR_6:.*]] = stablehlo.floor %[[ARG0]] : tensor --// CHECK: %[[COMPARE_15:.*]] = stablehlo.compare NE, %[[ARG0]], %[[FLOOR_6]] : (tensor, tensor) -> tensor --// CHECK: %[[COMPARE_16:.*]] = stablehlo.compare LT, %[[ARG0]], %[[CONSTANT_70]] : (tensor, tensor) -> tensor --// CHECK: %[[OR_0:.*]] = stablehlo.or %[[COMPARE_15]], %[[COMPARE_16]] : tensor -+// CHECK: %[[COMPARE_16:.*]] = stablehlo.compare NE, %[[ARG0]], %[[FLOOR_6]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_17:.*]] = stablehlo.compare LT, %[[ARG0]], %[[CONSTANT_70]] : (tensor, tensor) -> tensor -+// CHECK: %[[OR_0:.*]] = stablehlo.or %[[COMPARE_16]], %[[COMPARE_17]] : tensor - // CHECK: %[[CONSTANT_97:.*]] = stablehlo.constant dense<0x7FF8000000000000> : tensor - // CHECK: %[[SELECT_15:.*]] = stablehlo.select %[[OR_0]], %[[CONSTANT_97]], %[[SELECT_14]] : tensor, tensor - // CHECK: return %[[SELECT_15]] : tensor -@@ -3724,17 +3739,20 @@ - // CHECK: %[[DIVIDE_32:.*]] = stablehlo.divide %[[MULTIPLY_52]], %[[SINE_1]] : tensor - // CHECK: %[[SUBTRACT_18:.*]] = stablehlo.subtract %[[SUBTRACT_17]], %[[DIVIDE_32]] : tensor - // CHECK: %[[SELECT_12:.*]] = stablehlo.select %[[COMPARE_12]], %[[SUBTRACT_18]], %[[SUBTRACT_17]] : tensor, tensor --// CHECK: %[[COMPARE_13:.*]] = stablehlo.compare LE, %[[CONVERT_1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_13:.*]] = stablehlo.compare EQ, %[[CONVERT_1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_14:.*]] = stablehlo.compare LT, %[[CONVERT_1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor - // CHECK: %[[FLOOR_5:.*]] = stablehlo.floor %[[CONVERT_1]] : tensor --// CHECK: %[[COMPARE_14:.*]] = stablehlo.compare EQ, %[[CONVERT_1]], %[[FLOOR_5]] : (tensor, tensor) -> tensor --// CHECK: %[[AND_3:.*]] = stablehlo.and %[[COMPARE_13]], %[[COMPARE_14]] : tensor -+// CHECK: %[[COMPARE_15:.*]] = stablehlo.compare EQ, %[[CONVERT_1]], %[[FLOOR_5]] : (tensor, tensor) -> tensor -+// CHECK: %[[AND_3:.*]] = stablehlo.and %[[COMPARE_14]], %[[COMPARE_15]] : tensor - // CHECK: %[[CONSTANT_96:.*]] = stablehlo.constant dense<0x7FC00000> : tensor - // CHECK: %[[SELECT_13:.*]] = stablehlo.select %[[AND_3]], %[[CONSTANT_96]], %[[SELECT_12]] : tensor, tensor --// CHECK: %[[SELECT_14:.*]] = stablehlo.select %[[COMPARE_11]], %[[SELECT_13]], %[[MULTIPLY_42]] : tensor, tensor -+// CHECK: %[[CONSTANT_98:.*]] = stablehlo.constant dense<0xFF800000> : tensor -+// CHECK: %[[SELECT_16:.*]] = stablehlo.select %[[COMPARE_13]], %[[CONSTANT_98]], %[[SELECT_13]] : tensor, tensor -+// CHECK: %[[SELECT_14:.*]] = stablehlo.select %[[COMPARE_11]], %[[SELECT_16]], %[[MULTIPLY_42]] : tensor, tensor - // CHECK: %[[FLOOR_6:.*]] = stablehlo.floor %[[CONVERT_0]] : tensor --// CHECK: %[[COMPARE_15:.*]] = stablehlo.compare NE, %[[CONVERT_0]], %[[FLOOR_6]] : (tensor, tensor) -> tensor --// CHECK: %[[COMPARE_16:.*]] = stablehlo.compare LT, %[[CONVERT_0]], %[[CONSTANT_70]] : (tensor, tensor) -> tensor --// CHECK: %[[OR_0:.*]] = stablehlo.or %[[COMPARE_15]], %[[COMPARE_16]] : tensor -+// CHECK: %[[COMPARE_16:.*]] = stablehlo.compare NE, %[[CONVERT_0]], %[[FLOOR_6]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_17:.*]] = stablehlo.compare LT, %[[CONVERT_0]], %[[CONSTANT_70]] : (tensor, tensor) -> tensor -+// CHECK: %[[OR_0:.*]] = stablehlo.or %[[COMPARE_16]], %[[COMPARE_17]] : tensor - // CHECK: %[[CONSTANT_97:.*]] = stablehlo.constant dense<0x7FC00000> : tensor - // CHECK: %[[SELECT_15:.*]] = stablehlo.select %[[OR_0]], %[[CONSTANT_97]], %[[SELECT_14]] : tensor, tensor - // CHECK: %[[CONVERT_2:.*]] = stablehlo.convert %[[SELECT_15]] : (tensor) -> tensor -@@ -5282,6 +5300,32 @@ - - // ----- - -+// CHECK-LABEL: func.func @ragged_dot_mode_2_rank_lhs_lt_rhs( -+// CHECK-SAME: %[[ARG0:.*]]: tensor<2x3xf32>, -+// CHECK-SAME: %[[ARG1:.*]]: tensor<2x3x4xf32>, -+// CHECK-SAME: %[[ARG2:.*]]: tensor<2xi64>) -> tensor<2x3x3x4xf32> { -+// CHECK: %[[DOT_GENERAL_0:.*]] = stablehlo.dot_general %{{.*}}, %[[ARG1]], contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] : (tensor<2x3xf32>, tensor<2x3x4xf32>) -> tensor<3x3x4xf32> -+// CHECK: %[[DOT_GENERAL_1:.*]] = stablehlo.dot_general %{{.*}}, %[[ARG1]], contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] : (tensor<2x3xf32>, tensor<2x3x4xf32>) -> tensor<3x3x4xf32> -+// CHECK: %[[CONCATENATE_0:.*]] = stablehlo.concatenate{{.*}}dim = 0 : (tensor<1x3x3x4xf32>, tensor<1x3x3x4xf32>) -> tensor<2x3x3x4xf32> -+// CHECK: return %[[CONCATENATE_0]] : tensor<2x3x3x4xf32> -+// CHECK: } -+func.func @ragged_dot_mode_2_rank_lhs_lt_rhs(%lhs : tensor<2x3xf32>, %rhs : tensor<2x3x4xf32>, %group_sizes : tensor<2xi64>) -> tensor<2x3x3x4xf32> { -+ %0 = "chlo.ragged_dot"(%lhs, %rhs, %group_sizes) { -+ ragged_dot_dimension_numbers = #chlo.ragged_dot< -+ lhs_batching_dimensions = [], -+ rhs_batching_dimensions = [], -+ lhs_contracting_dimensions = [0], -+ rhs_contracting_dimensions = [0], -+ lhs_ragged_dimensions = [0], -+ rhs_group_dimensions = [] -+ >, -+ precision_config = [#chlo, #chlo] -+ } : (tensor<2x3xf32>, tensor<2x3x4xf32>, tensor<2xi64>) -> tensor<2x3x3x4xf32> -+ func.return %0 : tensor<2x3x3x4xf32> -+} -+ -+// ----- -+ - // CHECK-LABEL: func.func @ragged_dot_mode_3( - // CHECK-SAME: %[[ARG0:.*]]: tensor<2x3x5xf32>, - // CHECK-SAME: %[[ARG1:.*]]: tensor<2x5x7xf32>, -diff --ruN a/stablehlo/stablehlo/tests/ops_broadcasting.mlir b/stablehlo/stablehlo/tests/ops_broadcasting.mlir ---- stablehlo/stablehlo/tests/ops_broadcasting.mlir -+++ stablehlo/stablehlo/tests/ops_broadcasting.mlir -@@ -320,3 +320,14 @@ - return %0 : !stablehlo.token - } - -+// ----- -+ -+// Non-numpy explicit broadcast_dimensions: [3, 1] -> [3, 4, 5] -+// CHECK-LABEL: func @explicit_broadcast_dims -+func.func @explicit_broadcast_dims(%arg0: tensor<3x1xf64>) -> tensor<3x4x5xf64> { -+ // CHECK: %[[BCAST:.+]] = stablehlo.broadcast_in_dim %arg0, dims = [0, 1] : (tensor<3x1xf64>) -> tensor<3x4x5xf64> -+ // CHECK-NEXT: return %[[BCAST]] : tensor<3x4x5xf64> -+ %0 = "hlo_test_broadcast.broadcast_if_needed"(%arg0) {broadcast_dimensions = array} : (tensor<3x1xf64>) -> tensor<3x4x5xf64> -+ return %0 : tensor<3x4x5xf64> -+} -+ -diff --ruN a/stablehlo/stablehlo/tests/verify_convolution.mlir b/stablehlo/stablehlo/tests/verify_convolution.mlir ---- stablehlo/stablehlo/tests/verify_convolution.mlir -+++ stablehlo/stablehlo/tests/verify_convolution.mlir -@@ -78,6 +78,60 @@ - - // ----- - -+// CHECK-LABEL: func @convolution_mixed_fp8 -+func.func @convolution_mixed_fp8(%arg0 : tensor<100x26x26x32xf8E5M2>, -+ %arg1 : tensor<3x3x1x32xf8E4M3FN>) -> tensor<100x28x28x1xf8E5M2> { -+ %result = "stablehlo.convolution"(%arg0, %arg1) { -+ batch_group_count = 1 : i64, -+ dimension_numbers = #stablehlo.conv, -+ feature_group_count = 1 : i64, -+ lhs_dilation = array, -+ padding = dense<2> : tensor<2x2xi64>, -+ rhs_dilation = array, -+ window_strides = array -+ } : (tensor<100x26x26x32xf8E5M2>, tensor<3x3x1x32xf8E4M3FN>) -> -+ tensor<100x28x28x1xf8E5M2> -+ func.return %result : tensor<100x28x28x1xf8E5M2> -+} -+ -+// ----- -+ -+func.func @convolution_mismatched_element_types(%arg0: tensor<100x26x26x32xf32>, %arg1: tensor<3x3x1x32xf16>) -> tensor<100x28x28x1xf32> { -+ // expected-error@+1{{expects lhs and rhs to have compatible element type. Got: 'f32' and 'f16'}} -+ %result = "stablehlo.convolution"(%arg0, %arg1) { -+ batch_group_count = 1 : i64, -+ dimension_numbers = #stablehlo.conv, -+ feature_group_count = 1 : i64, -+ lhs_dilation = array, -+ padding = dense<2> : tensor<2x2xi64>, -+ rhs_dilation = array, -+ window_strides = array -+ } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf16>) -> tensor<100x28x28x1xf32> -+ func.return %result : tensor<100x28x28x1xf32> -+} -+ -+// ----- -+ - func.func @convolution(%arg0: tensor<2x2x3x4xf32>, %arg1: tensor<3x5x5x3xf32>) -> tensor<3x5x5x4xf32> { - // expected-error@+3{{Unexpected keyword stide}} - %0 = stablehlo.convolution(%arg0, %arg1) -diff --ruN a/stablehlo/stablehlo/tests/verify_dynamic_conv.mlir b/stablehlo/stablehlo/tests/verify_dynamic_conv.mlir ---- stablehlo/stablehlo/tests/verify_dynamic_conv.mlir -+++ stablehlo/stablehlo/tests/verify_dynamic_conv.mlir -@@ -48,6 +48,36 @@ - - // ----- - -+// CHECK-LABEL: func @dynamic_conv_mixed_fp8 -+func.func @dynamic_conv_mixed_fp8(%arg0 : tensor<100x26x26x32xf8E5M2>, -+ %arg1 : tensor<3x3x1x32xf8E4M3FN>) -> tensor<100x28x28x1xf8E5M2> { -+ %padding = stablehlo.constant dense<2> : tensor<2x2xi64> -+ %result = "stablehlo.dynamic_conv"(%arg0, %arg1, %padding) { -+ dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>, -+ feature_group_count = 1 : i64, -+ batch_group_count = 1 : i64 -+ } : (tensor<100x26x26x32xf8E5M2>, tensor<3x3x1x32xf8E4M3FN>, tensor<2x2xi64>) -> -+ tensor<100x28x28x1xf8E5M2> -+ func.return %result : tensor<100x28x28x1xf8E5M2> -+} -+ -+// ----- -+ -+func.func @dynamic_conv_mismatched_element_types(%arg0: tensor<100x26x26x32xf32>, -+ %arg1: tensor<3x3x1x32xf16>) -> tensor<100x28x28x1xf32> { -+ // expected-error@+2 {{expects lhs and rhs to have compatible element type. Got: 'f32' and 'f16'}} -+ %padding = stablehlo.constant dense<2> : tensor<2x2xi64> -+ %result = "stablehlo.dynamic_conv"(%arg0, %arg1, %padding) { -+ dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>, -+ feature_group_count = 1 : i64, -+ batch_group_count = 1 : i64 -+ } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf16>, tensor<2x2xi64>) -> -+ tensor<100x28x28x1xf32> -+ func.return %result : tensor<100x28x28x1xf32> -+} -+ -+// ----- -+ - func.func @dynamic_conv_c1(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x207xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+2 {{expects convolution arguments to have same number of dimensions. Got: 'tensor<1x8x8x207xf32>' and 'tensor<3x3x207xf32>'.}} -diff --ruN a/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_19_0.mlir b/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_19_0.mlir ---- stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_19_0.mlir -+++ stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_19_0.mlir -@@ -0,0 +1,56 @@ -+// RUN: stablehlo-opt --stablehlo-legalize-to-vhlo --vhlo-to-version='target=1.19.0' --verify-diagnostics --split-input-file %s -+ -+// expected-error @+1 {{failed to convert VHLO to v1.19.0}} -+module { -+ func.func @convolution_mixed_fp8(%arg0: tensor<100x26x26x32xf8E5M2>, %arg1: tensor<3x3x1x32xf8E4M3FN>) -> tensor<100x28x28x1xf8E5M2> { -+ // expected-error @+1 {{failed to legalize operation 'vhlo.convolution_v1' that was explicitly marked illegal}} -+ %result = "stablehlo.convolution"(%arg0, %arg1) { -+ batch_group_count = 1 : i64, -+ dimension_numbers = #stablehlo.conv, -+ feature_group_count = 1 : i64, -+ lhs_dilation = array, -+ padding = dense<2> : tensor<2x2xi64>, -+ rhs_dilation = array, -+ window_strides = array -+ } : (tensor<100x26x26x32xf8E5M2>, tensor<3x3x1x32xf8E4M3FN>) -> tensor<100x28x28x1xf8E5M2> -+ func.return %result : tensor<100x28x28x1xf8E5M2> -+ } -+} -+ -+// ----- -+ -+// expected-error @+1 {{failed to convert VHLO to v1.19.0}} -+module { -+ func.func @dynamic_conv_mixed_fp8(%arg0: tensor<100x26x26x32xf8E5M2>, %arg1: tensor<3x3x1x32xf8E4M3FN>, %arg2: tensor<2x2xi64>) -> tensor<100x28x28x1xf8E5M2> { -+ // expected-error @+1 {{failed to legalize operation 'vhlo.dynamic_conv_v2' that was explicitly marked illegal}} -+ %result = "stablehlo.dynamic_conv"(%arg0, %arg1, %arg2) { -+ batch_group_count = 1 : i64, -+ dimension_numbers = #stablehlo.conv, -+ feature_group_count = 1 : i64, -+ lhs_dilation = array, -+ rhs_dilation = array, -+ window_strides = array -+ } : (tensor<100x26x26x32xf8E5M2>, tensor<3x3x1x32xf8E4M3FN>, tensor<2x2xi64>) -> tensor<100x28x28x1xf8E5M2> -+ func.return %result : tensor<100x28x28x1xf8E5M2> -+ } -+} -diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp ---- stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp -+++ stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp -@@ -1668,18 +1668,27 @@ - digamma = mlir::stablehlo::SelectOp::create(rewriter, loc, needToReflect, - reflection, digamma); - -- // Digamma has poles at negative integers and zero; return nan for those. -- Value isLeZero = mlir::stablehlo::CompareOp::create( -- rewriter, loc, x, zero, mlir::stablehlo::ComparisonDirection::LE); -+ // Digamma has poles at negative integers and zero; return nan for -+ // negative integers, and -inf for zero. -+ Value isZero = mlir::stablehlo::CompareOp::create( -+ rewriter, loc, x, zero, mlir::stablehlo::ComparisonDirection::EQ); -+ Value isLtZero = mlir::stablehlo::CompareOp::create( -+ rewriter, loc, x, zero, mlir::stablehlo::ComparisonDirection::LT); - Value isInt = mlir::stablehlo::CompareOp::create( - rewriter, loc, x, mlir::stablehlo::FloorOp::create(rewriter, loc, x), - mlir::stablehlo::ComparisonDirection::EQ); -- Value isPole = mlir::stablehlo::AndOp::create(rewriter, loc, isLeZero, isInt); -- return mlir::stablehlo::SelectOp::create( -- rewriter, loc, isPole, -+ Value isNegativeInteger = -+ mlir::stablehlo::AndOp::create(rewriter, loc, isLtZero, isInt); -+ Value resultWithNan = mlir::stablehlo::SelectOp::create( -+ rewriter, loc, isNegativeInteger, - getConstantLike(rewriter, loc, std::numeric_limits::quiet_NaN(), - x), - digamma); -+ return mlir::stablehlo::SelectOp::create( -+ rewriter, loc, isZero, -+ getConstantLike(rewriter, loc, -std::numeric_limits::infinity(), -+ x), -+ resultWithNan); - } - - namespace { -@@ -2752,13 +2761,25 @@ - LogicalResult matchAndRewrite( - mlir::chlo::RaggedDotOp op, OpAdaptor, - ConversionPatternRewriter& rewriter) const override { -- if (op.getLhs().getType().getRank() < op.getRhs().getType().getRank()) { -+ chlo::RaggedDotDimensionNumbersAttr raggedDotDimensionNumbers = -+ op.getRaggedDotDimensionNumbers(); -+ ArrayRef lhsRaggedDimensions = -+ raggedDotDimensionNumbers.getLhsRaggedDimensions(); -+ if (lhsRaggedDimensions.empty()) { -+ return rewriter.notifyMatchFailure( -+ op, "lhs_ragged_dimensions must not be empty"); -+ } -+ const int64_t lhsRaggedDim = lhsRaggedDimensions[0]; -+ if (llvm::is_contained( -+ raggedDotDimensionNumbers.getLhsContractingDimensions(), -+ lhsRaggedDim)) { -+ return handleRaggedDotMode2(op, rewriter); -+ } else if (llvm::is_contained( -+ raggedDotDimensionNumbers.getLhsBatchingDimensions(), -+ lhsRaggedDim)) { -+ return handleRaggedDotMode3(op, rewriter); -+ } else { - return handleRaggedDotMode1(op, rewriter); -- } else if (op.getLhs().getType().getRank() < -- op.getResult().getType().getRank()) { -- return handleRaggedDotMode2(op, rewriter); -- } else { -- return handleRaggedDotMode3(op, rewriter); - } - } - }; -diff --ruN a/stablehlo/stablehlo/transforms/StablehloBroadcastLowering.cpp b/stablehlo/stablehlo/transforms/StablehloBroadcastLowering.cpp ---- stablehlo/stablehlo/transforms/StablehloBroadcastLowering.cpp -+++ stablehlo/stablehlo/transforms/StablehloBroadcastLowering.cpp -@@ -224,23 +224,51 @@ - mlir::RankedTensorType outputType = - getRankedTensorType(shape, inputType.getElementType()); - -- // Short circuit if no broadcasting is needed. -- if (inputType == outputType) return input; -- - int64_t inputRank = inputType.getRank(); - int64_t outputRank = outputType.getRank(); - if (inputRank > outputRank) - return emitError(loc, "input rank must be <= output rank, got ") - << inputRank << " vs " << outputRank; - -- size_t rankDiff = outputRank - inputRank; -+ // Construct broadcast dimensions (right-aligned for NumPy-style -+ // broadcasting). -+ auto broadcastDimensions = -+ llvm::to_vector(llvm::seq(outputRank - inputRank, outputRank)); -+ -+ return broadcastIfNeeded(builder, input, shape, broadcastDimensions); -+} -+ -+FailureOr broadcastIfNeeded(OpBuilder& builder, Value input, -+ const Dimensions& shape, -+ ArrayRef broadcastDimensions) { -+ LLVM_DEBUG(llvm::dbgs() << "[broadcastIfNeeded] Broadcasting input " -+ << input.getType() << " => " << toString(shape) -+ << "\n"); -+ auto loc = input.getLoc(); -+ mlir::RankedTensorType inputType = -+ dyn_cast(input.getType()); -+ if (!inputType) -+ return emitError(loc, "expected ranked tensor type for broadcast inputs"); -+ mlir::RankedTensorType outputType = -+ getRankedTensorType(shape, inputType.getElementType()); -+ -+ // Short circuit if no broadcasting is needed. -+ if (inputType == outputType) return input; -+ -+ int64_t inputRank = inputType.getRank(); -+ int64_t outputRank = outputType.getRank(); -+ if (inputRank > outputRank) -+ return emitError(loc, "input rank must be <= output rank, got ") -+ << inputRank << " vs " << outputRank; -+ -+ if (static_cast(broadcastDimensions.size()) != inputRank) -+ return emitError(loc, "broadcast_dimensions size (") -+ << broadcastDimensions.size() << ") must match input rank (" -+ << inputRank << ")"; -+ - auto inputShapeOrFail = getDimensions(input); - if (failed(inputShapeOrFail)) return failure(); - Dimensions inputShape = std::move(*inputShapeOrFail); -- -- // Construct broadcast dimensions. -- auto broadcastDimensions = -- llvm::to_vector(llvm::seq(outputRank - inputRank, outputRank)); - - // Construct the result type of the broadcast - // - If input is static and target shape is static, use static shape. -@@ -248,14 +276,30 @@ - // - If input is not bounded, but target shape is bounded, broadcast to - // the padded shape then call SetDimensionSize to make dynamic. - auto bcastShape = shape; -+ llvm::SmallVector isMapped(outputRank, false); -+ - for (int64_t i = 0; i < inputRank; ++i) { -+ int64_t resultIdx = broadcastDimensions[i]; -+ if (resultIdx < 0 || resultIdx >= outputRank) -+ return emitError(loc, "broadcast_dimensions index ") -+ << resultIdx << " out of bounds for output rank " << outputRank; -+ -+ isMapped[resultIdx] = true; -+ - int64_t inputDimSize = inputShape[i].size; -- int64_t resultIdx = i + rankDiff; - int64_t resultDimSize = shape[resultIdx].size; - if (inputDimSize != 1 && inputDimSize != resultDimSize) - return emitError(loc, "Cannot broadcast input: ") - << inputType << " to target shape " << toString(shape); - -+ if (inputShape[i].boundOp.has_value() && -+ !shape[resultIdx].boundOp.has_value()) { -+ return emitError( -+ loc, "cannot mix bounded and static dimensions in broadcast: ") -+ << "input dimension " << i << " is bounded, but target dimension " -+ << resultIdx << " is static"; -+ } -+ - if (!inputShape[i].boundOp.has_value() && - shape[resultIdx].boundOp.has_value()) { - // Use padded shape in broadcast. -@@ -263,17 +307,18 @@ - } - } - -- // Broadcast to padded size for remaining dimensions. -- for (size_t i = 0; i < rankDiff; ++i) { -- bcastShape[i] = DimensionInfo{shape[i].size}; -+ // Broadcast to padded size for remaining unmapped dimensions. -+ for (int64_t i = 0; i < outputRank; ++i) { -+ if (!isMapped[i]) { -+ bcastShape[i] = DimensionInfo{shape[i].size}; -+ } - } - - // Insert broadcast ops - mlir::RankedTensorType bcastType = - getRankedTensorType(bcastShape, inputType.getElementType()); -- LLVM_DEBUG( -- llvm::dbgs() << "[numpyBroadcastIfNeeded] Broadcast to padded type " -- << bcastType << "\n"); -+ LLVM_DEBUG(llvm::dbgs() << "[broadcastIfNeeded] Broadcast to padded type " -+ << bcastType << "\n"); - Value bcastOp = stablehlo::BroadcastInDimOp::create( - builder, loc, bcastType, input, broadcastDimensions); - if (bcastOp.getType() == outputType) return bcastOp; -diff --ruN a/stablehlo/stablehlo/transforms/StablehloBroadcastLowering.h b/stablehlo/stablehlo/transforms/StablehloBroadcastLowering.h ---- stablehlo/stablehlo/transforms/StablehloBroadcastLowering.h -+++ stablehlo/stablehlo/transforms/StablehloBroadcastLowering.h -@@ -70,6 +70,12 @@ - FailureOr numpyBroadcastIfNeeded(OpBuilder& builder, Value input, - const Dimensions& shape); - -+// Apply broadcasting to the given operand using the specified -+// broadcast_dimensions, returning an error if the operand is not broadcastable. -+FailureOr broadcastIfNeeded(OpBuilder& builder, Value input, -+ const Dimensions& shape, -+ ArrayRef broadcastDimensions); -+ - } // namespace stablehlo - } // namespace mlir - diff --git a/third_party/xla/third_party/stablehlo/workspace.bzl b/third_party/xla/third_party/stablehlo/workspace.bzl index 5f1d45c57e8e8b..c8508914e189ec 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 = "7b1b15781ccbd770f50c7eef4b0c3e03834649fd" - STABLEHLO_SHA256 = "d498ed1ba288eaf508369f76912e14d52d7f1f1a31ddaddd8ad4be40c8b7b4dc" + STABLEHLO_COMMIT = "639335932274930617d9724fb25ec9f2f52c5e40" + STABLEHLO_SHA256 = "dd590949798e26838d9ac3d6016def9b508caa60d0d7a69535ed2f5454ad30b8" # LINT.ThenChange(Google-internal path) tf_http_archive( diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc index 85852dba54c504..1db8afcf410d30 100644 --- a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc +++ b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc @@ -2481,8 +2481,8 @@ LogicalResult ExportXlaOp(BitcastConvertOp op, OpLoweringContext ctx) { LogicalResult ExportXlaOp(CollectiveBroadcastOp op, OpLoweringContext ctx) { auto& value_map = *ctx.values; - xla::XlaOp operand; - if (failed(GetXlaOp(op.getOperand(), value_map, &operand, op))) { + SmallVector operands; + if (failed(GetTuple(op.getOperation(), op.getOperands(), ctx, operands))) { return failure(); } auto replica_groups = Convert_replica_groups(op.getReplicaGroups(), op); @@ -2490,8 +2490,20 @@ LogicalResult ExportXlaOp(CollectiveBroadcastOp op, OpLoweringContext ctx) { return op.emitOpError(replica_groups.status().ToString()); } auto result = xla::CollectiveBroadcastWithDeviceList( - operand, **replica_groups, Convert_channel_handle(op.getChannelHandle())); - value_map[op->getResult(0)] = result; + operands, **replica_groups, Convert_channel_handle(op.getChannelHandle()), + op.getHasDynamicRoot()); + + // A collective_broadcast with more than one data operand produces a tuple. + mlir::FailureOr shape_or = + xla::ExtractXlaShape(op.getOperation()); + if (failed(shape_or)) { + return failure(); + } + if (shape_or->IsTuple()) { + BuildGetTupleElementsForTupleResults(op, result, ctx); + } else { + value_map[op->getResult(0)] = result; + } return success(); } @@ -3616,7 +3628,7 @@ LogicalResult ExportXlaOp(AsyncStartOp op, OpLoweringContext ctx) { xla::Shape input_shape = xla::ShapeUtil::MakeTupleShape( {xla::TypeToShape(op.getOperand(0).getType())}); - xla::Shape output_shape = xla::TypeToShape(collective_broadcast.getType()); + xla::Shape output_shape = xla::TypeToShape(collective_broadcast.getType(0)); xla::Shape start_shape = xla::ShapeUtil::MakeTupleShape({input_shape, output_shape}); (*ctx.values)[op.getResult()] = @@ -3698,7 +3710,7 @@ LogicalResult ExportXlaOp(AsyncDoneOp op, OpLoweringContext ctx) { (*ctx.values)[op.getResult()] = xla::internal::XlaBuilderFriend::BuildAsyncDone( ctx.builder, operand, - xla::TypeToShape(collective_broadcast.getType())); + xla::TypeToShape(collective_broadcast.getType(0))); return success(); } diff --git a/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.cc b/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.cc index 3fee8629835b1f..26ff56452437ea 100644 --- a/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.cc +++ b/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.cc @@ -1949,7 +1949,9 @@ void CollectiveBroadcastOp::build(OpBuilder& odsBuilder, } LogicalResult CollectiveBroadcastOp::verify() { - return hlo::verifyCollectiveBroadcastOp(getLoc(), getReplicaGroups()); + return hlo::verifyCollectiveBroadcastOp(getLoc(), (*this)->getOperands(), + getReplicaGroups(), + /*hasDynamicRoot=*/false); } //===----------------------------------------------------------------------===// diff --git a/third_party/xla/xla/python/version.h b/third_party/xla/xla/python/version.h index a2ae7e4250a98f..5f71f9ef3ecc56 100644 --- a/third_party/xla/xla/python/version.h +++ b/third_party/xla/xla/python/version.h @@ -18,6 +18,7 @@ limitations under the License. // An increasing version number to protect jax code against breaking changes. // In JAX, reference this via jax._src.lib.ifrt_version. -#define JAX_IFRT_VERSION_NUMBER 67 // Scan dimension in TPU MLIR dialect. +#define JAX_IFRT_VERSION_NUMBER \ + 68 // StableHlo CollectiveBroadcastOp Compatibility #endif // XLA_PYTHON_VERSION_H_