From 24581658029f6a9e2b4ded6b80b67bf6f616ed8b Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Mon, 24 Aug 2026 15:17:48 -0700 Subject: [PATCH 01/28] Validate input shapes in the LSTM gradient kernels --- tensorflow/core/kernels/rnn/lstm_ops.cc | 75 ++++++++++++++++ .../kernel_tests/nn_ops/rnn_cell_test.py | 87 +++++++++++++++++++ 2 files changed, 162 insertions(+) diff --git a/tensorflow/core/kernels/rnn/lstm_ops.cc b/tensorflow/core/kernels/rnn/lstm_ops.cc index 3609541988cccd..034acb87f08516 100644 --- a/tensorflow/core/kernels/rnn/lstm_ops.cc +++ b/tensorflow/core/kernels/rnn/lstm_ops.cc @@ -593,6 +593,35 @@ class LSTMBlockCellGradOp : public OpKernel { const Tensor* h_grad_tensor = nullptr; OP_REQUIRES_OK(ctx, ctx->input("h_grad", &h_grad_tensor)); + // The dimensions read below, and the sanity checks that follow them, index + // dimension 1 of these tensors. That is a fatal check rather than an error + // when a tensor has a lower rank, so validate the ranks first, the way + // LSTMBlockCellOp does before using the same inputs. + const auto check_rank = [](const Tensor* t, const char* name, + int expected) -> absl::Status { + if (t->dims() != expected) { + return absl::InvalidArgumentError(absl::StrCat( + name, " must be rank ", expected, " but is rank ", t->dims())); + } + return absl::OkStatus(); + }; + OP_REQUIRES_OK(ctx, check_rank(x_tensor, "x", 2)); + OP_REQUIRES_OK(ctx, check_rank(cs_prev_tensor, "cs_prev", 2)); + OP_REQUIRES_OK(ctx, check_rank(h_prev_tensor, "h_prev", 2)); + OP_REQUIRES_OK(ctx, check_rank(w_tensor, "w", 2)); + OP_REQUIRES_OK(ctx, check_rank(wci_tensor, "wci", 1)); + OP_REQUIRES_OK(ctx, check_rank(wcf_tensor, "wcf", 1)); + OP_REQUIRES_OK(ctx, check_rank(wco_tensor, "wco", 1)); + OP_REQUIRES_OK(ctx, check_rank(b_tensor, "b", 1)); + OP_REQUIRES_OK(ctx, check_rank(i_tensor, "i", 2)); + OP_REQUIRES_OK(ctx, check_rank(cs_tensor, "cs", 2)); + OP_REQUIRES_OK(ctx, check_rank(f_tensor, "f", 2)); + OP_REQUIRES_OK(ctx, check_rank(o_tensor, "o", 2)); + OP_REQUIRES_OK(ctx, check_rank(ci_tensor, "ci", 2)); + OP_REQUIRES_OK(ctx, check_rank(co_tensor, "co", 2)); + OP_REQUIRES_OK(ctx, check_rank(cs_grad_tensor, "cs_grad", 2)); + OP_REQUIRES_OK(ctx, check_rank(h_grad_tensor, "h_grad", 2)); + const int64_t batch_size = x_tensor->dim_size(0); const int64_t input_size = x_tensor->dim_size(1); const int64_t cell_size = cs_prev_tensor->dim_size(1); @@ -1295,6 +1324,52 @@ class BlockLSTMGradOp : public OpKernel { const Tensor* h_grad = nullptr; OP_REQUIRES_OK(ctx, ctx->input("h_grad", &h_grad)); + OP_REQUIRES(ctx, cs_grad->dims() == 3, + absl::InvalidArgumentError(absl::StrCat( + "cs_grad must be rank 3. Received: ", cs_grad->dims()))); + OP_REQUIRES(ctx, h_grad->dims() == 3, + absl::InvalidArgumentError(absl::StrCat( + "h_grad must be rank 3. Received: ", h_grad->dims()))); + + // `timelen`, `batch_size` and `cell_size` are taken from x and w, and every + // other input is sliced with them below. A tensor whose dimensions + // disagree is read out of bounds rather than failing the op: a zero batch + // on one side and a non-zero batch on the other crashes, and a shorter + // time dimension trips a fatal check while slicing. LSTMBlockCellGradOp + // already validates its own inputs this way. + const auto check_dims = [&](const Tensor* t, const char* name, + bool has_time) -> absl::Status { + const int batch_dim = has_time ? 1 : 0; + if (has_time && t->dim_size(0) != timelen) { + return absl::InvalidArgumentError( + absl::StrCat(name, ".dim_size(0) != timelen: ", t->dim_size(0), + " vs. ", timelen)); + } + if (t->dim_size(batch_dim) != batch_size) { + return absl::InvalidArgumentError(absl::StrCat( + name, ".dim_size(", batch_dim, + ") != batch_size: ", t->dim_size(batch_dim), " vs. ", batch_size)); + } + if (t->dim_size(batch_dim + 1) != cell_size) { + return absl::InvalidArgumentError( + absl::StrCat(name, ".dim_size(", batch_dim + 1, + ") != cell_size: ", t->dim_size(batch_dim + 1), + " vs. ", cell_size)); + } + return absl::OkStatus(); + }; + OP_REQUIRES_OK(ctx, check_dims(cs_prev_tensor, "cs_prev", false)); + OP_REQUIRES_OK(ctx, check_dims(h_prev_tensor, "h_prev", false)); + OP_REQUIRES_OK(ctx, check_dims(i_out, "i", true)); + OP_REQUIRES_OK(ctx, check_dims(cs_out, "cs", true)); + OP_REQUIRES_OK(ctx, check_dims(f_out, "f", true)); + OP_REQUIRES_OK(ctx, check_dims(o_out, "o", true)); + OP_REQUIRES_OK(ctx, check_dims(ci_out, "ci", true)); + OP_REQUIRES_OK(ctx, check_dims(co_out, "co", true)); + OP_REQUIRES_OK(ctx, check_dims(h_out, "h", true)); + OP_REQUIRES_OK(ctx, check_dims(cs_grad, "cs_grad", true)); + OP_REQUIRES_OK(ctx, check_dims(h_grad, "h_grad", true)); + TensorShape batch_input_shape({timelen, batch_size, input_size}); Tensor* x_grad; OP_REQUIRES_OK(ctx, diff --git a/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py b/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py index a1f4d741ba2782..fd3229bdb6fc0d 100644 --- a/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py @@ -1468,6 +1468,93 @@ def testLSTMBlockInvalidArgument(self): ) ) + def _blockLSTMGradInputs(self, timelen, batch_size, input_size, cell_size, + x_batch_size=None): + """Returns a consistent set of BlockLSTMGrad inputs, x batch overridable.""" + if x_batch_size is None: + x_batch_size = batch_size + def rand(*shape): + return random_ops.random_uniform(list(shape), dtype=dtypes.float32) + + def per_step(): + return rand(timelen, batch_size, cell_size) + + return dict( + seq_len_max=constant_op.constant(timelen, dtype=dtypes.int64), + x=rand(timelen, x_batch_size, input_size), + cs_prev=rand(batch_size, cell_size), + h_prev=rand(batch_size, cell_size), + w=rand(input_size + cell_size, 4 * cell_size), + wci=rand(cell_size), + wcf=rand(cell_size), + wco=rand(cell_size), + b=rand(4 * cell_size), + i=per_step(), + cs=per_step(), + f=per_step(), + o=per_step(), + ci=per_step(), + co=per_step(), + h=per_step(), + cs_grad=per_step(), + h_grad=per_step(), + use_peephole=True) + + def testBlockLSTMGradBatchSizeMismatch(self): + # Test case for GitHub issue 112746. x carried a different batch size from + # the other inputs, which were then sliced with the batch size taken from + # x and read out of bounds, crashing the process. + kwargs = self._blockLSTMGradInputs(3, 2, 4, 5, x_batch_size=0) + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, + "batch_size"): + self.evaluate(gen_rnn_ops.BlockLSTMGrad(**kwargs)) + + def testBlockLSTMGradTimeLenMismatch(self): + # A shorter time dimension on one of the per step inputs was sliced past + # its end. + kwargs = self._blockLSTMGradInputs(3, 2, 4, 5) + kwargs["i"] = random_ops.random_uniform([1, 2, 5], dtype=dtypes.float32) + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, "timelen"): + self.evaluate(gen_rnn_ops.BlockLSTMGrad(**kwargs)) + + def testBlockLSTMGradConsistentShapesStillWork(self): + # Consistent shapes must keep working, including a zero batch size. + for batch_size in (0, 2): + kwargs = self._blockLSTMGradInputs(3, batch_size, 4, 5) + self.evaluate(gen_rnn_ops.BlockLSTMGrad(**kwargs)) + + def testLSTMBlockCellGradInvalidRank(self): + # Test case for GitHub issue 112746. A rank 1 x was indexed at dimension 1 + # before any shape check ran, which aborted the process. + batch_size, input_size, cell_size = 2, 4, 5 + def rand(*shape): + return random_ops.random_uniform(list(shape), dtype=dtypes.float32) + + def per_step(): + return rand(batch_size, cell_size) + + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, + "x must be rank 2"): + self.evaluate( + gen_rnn_ops.LSTMBlockCellGrad( + x=constant_op.constant([], dtype=dtypes.float32), + cs_prev=per_step(), + h_prev=per_step(), + w=rand(input_size + cell_size, 4 * cell_size), + wci=rand(cell_size), + wcf=rand(cell_size), + wco=rand(cell_size), + b=rand(4 * cell_size), + i=per_step(), + cs=per_step(), + f=per_step(), + o=per_step(), + ci=per_step(), + co=per_step(), + cs_grad=per_step(), + h_grad=per_step(), + use_peephole=True)) + class BidirectionalRNNTest(test.TestCase): From 9278060fe7482feae0612f230e8f9ea78a65bdc5 Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Mon, 24 Aug 2026 15:40:57 -0700 Subject: [PATCH 02/28] Check the rank inside the BlockLSTMGrad dimension helper --- tensorflow/core/kernels/rnn/lstm_ops.cc | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tensorflow/core/kernels/rnn/lstm_ops.cc b/tensorflow/core/kernels/rnn/lstm_ops.cc index 034acb87f08516..9da23a5bc63b50 100644 --- a/tensorflow/core/kernels/rnn/lstm_ops.cc +++ b/tensorflow/core/kernels/rnn/lstm_ops.cc @@ -1324,21 +1324,21 @@ class BlockLSTMGradOp : public OpKernel { const Tensor* h_grad = nullptr; OP_REQUIRES_OK(ctx, ctx->input("h_grad", &h_grad)); - OP_REQUIRES(ctx, cs_grad->dims() == 3, - absl::InvalidArgumentError(absl::StrCat( - "cs_grad must be rank 3. Received: ", cs_grad->dims()))); - OP_REQUIRES(ctx, h_grad->dims() == 3, - absl::InvalidArgumentError(absl::StrCat( - "h_grad must be rank 3. Received: ", h_grad->dims()))); - // `timelen`, `batch_size` and `cell_size` are taken from x and w, and every // other input is sliced with them below. A tensor whose dimensions // disagree is read out of bounds rather than failing the op: a zero batch // on one side and a non-zero batch on the other crashes, and a shorter // time dimension trips a fatal check while slicing. LSTMBlockCellGradOp - // already validates its own inputs this way. + // already validates its own inputs this way. The rank is checked here as + // well so that the dimension lookups below are valid for any caller. const auto check_dims = [&](const Tensor* t, const char* name, bool has_time) -> absl::Status { + const int expected_dims = has_time ? 3 : 2; + if (t->dims() != expected_dims) { + return absl::InvalidArgumentError( + absl::StrCat(name, " must be rank ", expected_dims, + " but is rank ", t->dims())); + } const int batch_dim = has_time ? 1 : 0; if (has_time && t->dim_size(0) != timelen) { return absl::InvalidArgumentError( From a9c8bae9067fd918905d0ee0477d8f23f5cd8ef9 Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Tue, 25 Aug 2026 23:35:42 -0700 Subject: [PATCH 03/28] Accept the shape function error in the rank validation test --- tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py b/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py index fd3229bdb6fc0d..3602c7a46479c2 100644 --- a/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py @@ -1525,7 +1525,10 @@ def testBlockLSTMGradConsistentShapesStillWork(self): def testLSTMBlockCellGradInvalidRank(self): # Test case for GitHub issue 112746. A rank 1 x was indexed at dimension 1 - # before any shape check ran, which aborted the process. + # before any shape check ran, which aborted the process. Eager execution + # reaches the kernel check ("x must be rank 2"), while graph construction + # is stopped earlier by the op's shape function ("Shape must be rank 2"), + # which surfaces as a ValueError, so both are accepted. batch_size, input_size, cell_size = 2, 4, 5 def rand(*shape): return random_ops.random_uniform(list(shape), dtype=dtypes.float32) @@ -1533,8 +1536,8 @@ def rand(*shape): def per_step(): return rand(batch_size, cell_size) - with self.assertRaisesRegex(errors_impl.InvalidArgumentError, - "x must be rank 2"): + with self.assertRaisesRegex((ValueError, errors_impl.InvalidArgumentError), + "must be rank 2"): self.evaluate( gen_rnn_ops.LSTMBlockCellGrad( x=constant_op.constant([], dtype=dtypes.float32), From 22f62b5c3b1989c5c28e7a94555d0eee82005827 Mon Sep 17 00:00:00 2001 From: Prasanna Dabi <65734642+prasanna8585@users.noreply.github.com> Date: Wed, 23 Sep 2026 06:06:42 -0700 Subject: [PATCH 04/28] PR #46773: Fix unguarded int overflow in ShardingParam and validate on deserialize Imported from GitHub PR https://github.com/openxla/xla/pull/46773 ShardingParam::MinorToMajor::ToDeviceList computed a cumulative axis size product using 32-bit int with no overflow protection, and ShardingParam::verify() computed a dimension-shard product the same unguarded way before using it as a modulo divisor -- if the product wraps to exactly 0, that is a division by zero.\n\nConfirmed via a faithful standalone extraction of both arithmetic patterns, compiled with UBSan+ASan: both overflow with individually plausible input values, and the verify() case produces a real SIGFPE crash when the product wraps to 0.\n\nSeparately, ShardingParam::FromProto never called verify() at all, so a deserialized ShardingParam was never validated before its axis_sizes/dim_shards were used downstream.\n\nThis adds an overflow-safe int64_t accumulation with an explicit range check in verify(), calls verify() from FromProto, and hardens ToDeviceList to fail closed rather than silently propagate a wrapped value. Copybara import of the project: -- c015dcb95f89e46df503008a3de0333867b60dd7 by prasanna8585 <65734642+prasanna8585@users.noreply.github.com>: Fix unguarded int overflow in ShardingParam and validate on deserialize ShardingParam::MinorToMajor::ToDeviceList computed a cumulative axis size product using 32-bit int with no overflow protection (cum_size *= size), and ShardingParam::verify() computed a dimension- shard product the same unguarded way, then used it as a modulo divisor (NumDevices() % total_dim_shards_size) -- if the product wraps to exactly 0, this is a division by zero. Confirmed via a faithful standalone extraction of both arithmetic patterns, compiled with UBSan+ASan: both overflow with individually plausible input values, and the verify() case produces a real SIGFPE crash when the product wraps to 0. Separately, ShardingParam::FromProto never called verify() at all, so a deserialized ShardingParam (e.g. via xla/python/ifrt_proxy/server/ifrt_backend.cc's handling of MakeArraysFromHostBufferShards/Assemble requests) was never validated before its axis_sizes/dim_shards were used downstream. This adds an overflow-safe int64_t accumulation with an explicit range check in verify(), calls verify() from FromProto so deserialized input is actually validated, and hardens ToDeviceList to fail closed (CHECK) rather than silently propagate a wrapped value into device- index computation, on the assumption that verify() should have already caught it upstream. -- 961720e9258f1e54e674d12ecc821119f459cbf3 by prasanna8585 <65734642+prasanna8585@users.noreply.github.com>: Address review: move overflow check to verify(), use __builtin_mul_overflow, reject dim_shard<=0 - MinorToMajor::verify() now validates the axis_sizes cumulative product before it reaches ToDeviceList(), per @hyeontaek: validation should happen close to where unvalidated data enters, not repeated downstream. - Both overflow checks now use __builtin_mul_overflow instead of a post-hoc numeric_limits::max() comparison, matching the pattern in xla/python/ifrt/shape.cc, so correctness doesn't depend on sizeof(int) < sizeof(int64_t). - ShardingParam::verify() now rejects dim_shard <= 0 up front, so the separate zero-divisor check before the modulo is no longer needed. -- 75a4bb962c7687082c65af78b290f4103f39bfb3 by prasanna8585 <65734642+prasanna8585@users.noreply.github.com>: Add unit test for ShardingParam::verify() overflow rejection Covers axis_sizes product overflow, dim_shards product overflow, non-positive dim_shard rejection, and the valid-input path, exercising ShardingParam::verify() directly rather than going through FromProto/deserialization. -- f038531ededf56b60638e9526a04117929169486 by prasanna8585 : Address review: use int directly, trim comments, drop redundant check hyeontaek's review, all five points: 1-3. cum_size and total_dim_shards_size are now plain int rather than int64_t. __builtin_mul_overflow checks overflow against its output argument's type directly, so once the output is int there is no need to widen the inputs to int64_t and separately compare against std::numeric_limits::max() -- the builtin already does exactly that check when the output type is int. Verified directly: an isolated int-only build of this exact pattern correctly flags 2 * 1,500,000,000 (the real regression test's case) and the precise 46340/46341 boundary around INT_MAX, and correctly does not flag values within range. Removed the now-unused include from this file (the test file's own include is untouched, out of scope here). 2. Trimmed the three comments this PR added to state the same thing with less text. 4. sharding_param_test.cc: EXPECT_TRUE(sharding.verify().ok()) -> EXPECT_OK(sharding.verify()), matching the existing absl/status/status_matchers.h include already used by the rest of this file's assertions. 5. Removed the redundant overflow check in ToDeviceList(). Checked its only two real callers (sharding_spec.cc, sharding.cc) before removing it: both operate on an already-constructed ShardingParam member, consistent with the original commit's own stated design ("validate on deserialize") and with verify() being the sole enforcement point. cum_size is now plain int throughout, updated in the same place it's pushed, with no separate overflow check -- trusting the invariant verify() already established, per the review's own reasoning for why the duplicate check adds nothing. -- ef7d14e8662c4d4d8207890023cbeb66d67f205a by prasanna8585 : Address review: remove two comments explaining assumptions established elsewhere dimitar-asenov's feedback: 1. "Already rejected by verify() at deserialization time." above the cum_size *= size multiplication in ToDeviceList() was confusing -- the reviewer's own question ("What is rejected?") is the problem: the comment references an invariant enforced by a verify() call this function doesn't itself make, doesn't say which of the file's two verify() functions, and gives no way to confirm from this function's own scope that the invariant actually holds for every caller. Removed rather than clarified, since the underlying claim isn't something ToDeviceList() can locally establish or enforce. 2. "`total_dim_shards_size` is positive here, so no zero-divisor check is needed below." above the NumDevices() % total_dim_shards_size check in verify() -- explicit ask: delete it. As the reviewer put it, divisors not being zero is an assumption we don't generally need to explain, and if the code doesn't actually assure that on its own, that's a separate issue from this comment. No logic changed in either case -- both are comment-only deletions. -- f9ab8c868ba4a02a9d4f9d223dc9844f0a3c00b8 by prasanna8585 : Fix build: add missing tsl/platform/errors.h include thcmbs reported a integration failure: sharding_param_test.cc used tsl::error::INVALID_ARGUMENT (via the StatusIs matcher) without including the header that declares the tsl::error namespace, so it only resolved by chance in the environment this test was originally written against. hyeontaek's fix: add #include "xla/tsl/platform/errors.h". Also added the corresponding //xla/tsl/platform:errors Bazel dependency to the sharding_param_test target in BUILD -- the include alone isn't sufficient without it. Confirmed both the include ordering (alphabetical within the absl/llvm/xla group, matching this file's existing style) and the BUILD dep ordering (// targets before @ targets, matching every other target in this same file) against the existing conventions already used elsewhere in this same BUILD file and by sibling files (layout_test.cc, serdes.cc, serdes_test.cc) that already depend on this same header. Merging this change closes #46773 PiperOrigin-RevId: 986704240 --- third_party/xla/xla/python/ifrt/ir/BUILD | 12 +++ .../xla/xla/python/ifrt/ir/sharding_param.cc | 28 ++++++- .../xla/python/ifrt/ir/sharding_param_test.cc | 83 +++++++++++++++++++ 3 files changed, 120 insertions(+), 3 deletions(-) create mode 100644 third_party/xla/xla/python/ifrt/ir/sharding_param_test.cc diff --git a/third_party/xla/xla/python/ifrt/ir/BUILD b/third_party/xla/xla/python/ifrt/ir/BUILD index 4ced1e754822df..8a69d0b848b6d7 100644 --- a/third_party/xla/xla/python/ifrt/ir/BUILD +++ b/third_party/xla/xla/python/ifrt/ir/BUILD @@ -150,6 +150,18 @@ cc_library( ], ) +xla_cc_test( + name = "sharding_param_test", + srcs = ["sharding_param_test.cc"], + deps = [ + ":sharding_param", + "//xla/tsl/platform:errors", + "@com_google_absl//absl/status:status_matchers", + "@com_google_googletest//:gtest_main", + "@llvm-project//llvm:Support", + ], +) + tf_proto_library( name = "sharding_param_proto", srcs = ["sharding_param.proto"], diff --git a/third_party/xla/xla/python/ifrt/ir/sharding_param.cc b/third_party/xla/xla/python/ifrt/ir/sharding_param.cc index 4b7fb83df7389d..34b683ddf9d72c 100644 --- a/third_party/xla/xla/python/ifrt/ir/sharding_param.cc +++ b/third_party/xla/xla/python/ifrt/ir/sharding_param.cc @@ -151,6 +151,16 @@ absl::Status ShardingParam::MinorToMajor::verify() const { absl::StrJoin(axis_sizes, "x"))); } } + // The cumulative product of `axis_sizes` must fit in an `int`, since + // ToDeviceList() narrows to it. + int cum_size = 1; + for (const int size : axis_sizes) { + if (__builtin_mul_overflow(cum_size, size, &cum_size)) { + return absl::InvalidArgumentError( + absl::StrCat("`axis_sizes` product overflows: [", + absl::StrJoin(axis_sizes, "x"), "]")); + } + } return absl::OkStatus(); } @@ -248,7 +258,17 @@ absl::Status ShardingParam::verify() const { // distributed over the device mesh. int total_dim_shards_size = 1; for (const int dim_shard : dim_shards()) { - total_dim_shards_size *= dim_shard; + if (dim_shard <= 0) { + return absl::InvalidArgumentError( + absl::StrCat("`dim_shards` must be positive, got ", dim_shard, + " in [", absl::StrJoin(dim_shards(), "x"), "]")); + } + if (__builtin_mul_overflow(total_dim_shards_size, dim_shard, + &total_dim_shards_size)) { + return absl::InvalidArgumentError( + absl::StrCat("`dim_shards` product overflows: ", + absl::StrJoin(dim_shards(), "x"))); + } } if (NumDevices() % total_dim_shards_size != 0) { return absl::InvalidArgumentError(absl::StrCat( @@ -378,8 +398,10 @@ absl::StatusOr ShardingParam::FromProto( unreduced_axes = std::vector(proto.unreduced_axes().begin(), proto.unreduced_axes().end()); } - return ShardingParam(std::move(dim_shards), std::move(minor_to_major), - std::move(unreduced_axes)); + ShardingParam param(std::move(dim_shards), std::move(minor_to_major), + std::move(unreduced_axes)); + ABSL_RETURN_IF_ERROR(param.verify()); + return param; } absl::Status ShardingParam::ToProto(ShardingParamProto& proto, diff --git a/third_party/xla/xla/python/ifrt/ir/sharding_param_test.cc b/third_party/xla/xla/python/ifrt/ir/sharding_param_test.cc new file mode 100644 index 00000000000000..9b698d7aa5f75f --- /dev/null +++ b/third_party/xla/xla/python/ifrt/ir/sharding_param_test.cc @@ -0,0 +1,83 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/python/ifrt/ir/sharding_param.h" + +#include +#include + +#include +#include "absl/status/status_matchers.h" +#include "llvm/ADT/SmallVector.h" +#include "xla/tsl/platform/errors.h" + +namespace xla { +namespace ifrt { +namespace { + +using ::absl_testing::StatusIs; +using ::testing::HasSubstr; + +TEST(ShardingParamTest, AxisSizesProductOverflowIsRejected) { + // Two axis sizes whose product exceeds std::numeric_limits::max() + // (2 * 1,500,000,000 > INT_MAX), individually representable as plain ints. + ShardingParam::MinorToMajor minor_to_major; + minor_to_major.permutation = {0, 1}; + minor_to_major.axis_sizes = {2, 1500000000}; + + ShardingParam sharding(/*dim_shards=*/{1, 1}, minor_to_major); + + EXPECT_THAT(sharding.verify(), + StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("`axis_sizes` product overflows"))); +} + +TEST(ShardingParamTest, DimShardsProductOverflowIsRejected) { + ShardingParam::MinorToMajor minor_to_major; + minor_to_major.permutation = {0}; + minor_to_major.axis_sizes = {4}; + + ShardingParam sharding(/*dim_shards=*/{2, 1500000000}, minor_to_major); + + EXPECT_THAT(sharding.verify(), + StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("`dim_shards` product overflows"))); +} + +TEST(ShardingParamTest, NonPositiveDimShardIsRejected) { + ShardingParam::MinorToMajor minor_to_major; + minor_to_major.permutation = {0}; + minor_to_major.axis_sizes = {4}; + + ShardingParam sharding(/*dim_shards=*/{0}, minor_to_major); + + EXPECT_THAT(sharding.verify(), + StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("`dim_shards` must be positive"))); +} + +TEST(ShardingParamTest, ValidShardingIsAccepted) { + ShardingParam::MinorToMajor minor_to_major; + minor_to_major.permutation = {0, 1}; + minor_to_major.axis_sizes = {2, 4}; + + ShardingParam sharding(/*dim_shards=*/{2, 4}, minor_to_major); + + EXPECT_OK(sharding.verify()); +} + +} // namespace +} // namespace ifrt +} // namespace xla From 8b78b2caa6b1a6cdf0ec1d93b32cb3dcad1b2e09 Mon Sep 17 00:00:00 2001 From: dlcompilers-infra-bot Date: Wed, 23 Sep 2026 06:18:48 -0700 Subject: [PATCH 05/28] PR #49012: [XLA:FFI] Add RegisterWindow + GetWindow to the collectives FFI extension MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/49012 πŸ“ Summary of Changes Adds a collective-memory window API to the backend-agnostic collectives FFI extension: `register_window` (Prepare) registers a batch of already-allocated buffers with the clique, and `get_window` (Init/Execute) returns an opaque `XLA_FFI_Window*` per buffer along with the byte offset of the queried address within that window. Provides a C++ wrapper (`ffi::Communicator::RegisterWindow` / `GetWindow` returning `WindowLookup{window, offset}`). Bumps the extension minor version to 2. 🎯 Justification Backend-agnostic API surface for external custom-call handlers to reuse XLA-managed collective memory windows instead of allocating and registering their own. The offset is what handlers pass to the backend's window accessors (e.g. `ncclGetPeerDevicePointer(win, offset, peer, &out)`). Split out from #48557 per review request; GPU wiring + tests remain in #48557 and will merge after this. πŸš€ Kind of Contribution ✨ New Feature πŸ“Š Benchmark (for Performance Improvements) N/A πŸ§ͺ Unit Tests: N/A in this PR β€” window API is exercised by `CollectiveOpsTestFFI.PublicApiWindow` in the GPU-wiring follow-up (#48557). πŸ§ͺ Execution Tests: N/A Copybara import of the project: -- f4a4dd69232de42f2049e0f30a1b9fa24f732ba2 by Sevin F. Varoglu : [XLA:FFI] Add RegisterWindow + GetWindow to the collectives FFI extension -- c35f0ded55c1e24354cca4aa7261f04b5e2c1958 by Sevin F. Varoglu : Rename WindowRegister and add unit tests -- 73bffbd02c09682fde3b1293773f79a67580a88a by Sevin F. Varoglu : Remove SupportsWindow Merging this change closes #49012 PiperOrigin-RevId: 986709932 --- third_party/xla/xla/ffi/api/collectives_api.h | 68 +++++++++++ .../xla/xla/ffi/api/collectives_c_api.h | 99 ++++++++++++++-- .../xla/xla/ffi/api/collectives_ffi_test.cc | 111 ++++++++++++++++- .../xla/xla/ffi/collectives_ffi_test.cc | 112 +++++++++++++++++- 4 files changed, 376 insertions(+), 14 deletions(-) diff --git a/third_party/xla/xla/ffi/api/collectives_api.h b/third_party/xla/xla/ffi/api/collectives_api.h index 4d832aa9616781..c76f4686f34863 100644 --- a/third_party/xla/xla/ffi/api/collectives_api.h +++ b/third_party/xla/xla/ffi/api/collectives_api.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef XLA_FFI_API_COLLECTIVES_API_H_ #define XLA_FFI_API_COLLECTIVES_API_H_ +#include #include #include @@ -33,6 +34,17 @@ enum class GroupMode { kFlattenedId = XLA_FFI_GROUP_FLATTENED_ID, }; +struct CollectiveMemoryRegion { + const void* buffer; + size_t byte_size; + uint64_t flags = 0; +}; + +struct WindowLookup { + XLA_FFI_Window* window; + size_t offset; +}; + namespace internal { // C++ wrapper for the XLA FFI Collectives extension API. @@ -87,6 +99,62 @@ class CommunicatorContextBase { return args.communicator; } + //===--------------------------------------------------------------------===// + // Collective memory window + //===--------------------------------------------------------------------===// + // + // Request window registration for a batch of already-allocated buffers in + // Prepare; look up an opaque backend-defined window handle per buffer in + // Init/Execute. The handler reinterprets the window and calls the backend's + // collective device APIs directly to obtain local, peer, and multicast + // pointers. Allocation stays on the JAX-side. + + Status RequestWindow(GroupMode group_mode, + const std::vector>& groups, + int64_t communication_id, + const std::vector& regions) { + std::vector raw_groups = ToRawGroups(groups); + std::vector raw_regions; + raw_regions.reserve(regions.size()); + for (const CollectiveMemoryRegion& r : regions) { + raw_regions.push_back( + XLA_FFI_CollectiveMemoryRegion{r.buffer, r.byte_size, r.flags}); + } + XLA_FFI_Window_Request_Args args; + args.struct_size = XLA_FFI_Window_Request_Args_STRUCT_SIZE; + args.extension_start = nullptr; + args.group_mode = static_cast(group_mode); + args.groups = raw_groups.data(); + args.num_groups = raw_groups.size(); + args.communication_id = communication_id; + args.regions = raw_regions.data(); + args.num_regions = raw_regions.size(); + if (XLA_FFI_Error* err = ext_->request_window(ext_, &args)) { + return ErrorPolicy::TakeError(api_, err); + } + return ErrorPolicy::Ok(); + } + + StatusOr GetWindow( + GroupMode group_mode, const std::vector>& groups, + int64_t communication_id, const void* buffer) { + std::vector raw_groups = ToRawGroups(groups); + XLA_FFI_Window_Get_Args args; + args.struct_size = XLA_FFI_Window_Get_Args_STRUCT_SIZE; + args.extension_start = nullptr; + args.group_mode = static_cast(group_mode); + args.groups = raw_groups.data(); + args.num_groups = raw_groups.size(); + args.communication_id = communication_id; + args.buffer = buffer; + args.window = nullptr; + args.window_offset = 0; + if (XLA_FFI_Error* err = ext_->get_window(ext_, &args)) { + return StatusOr(ErrorPolicy::TakeError(api_, err)); + } + return WindowLookup{args.window, args.window_offset}; + } + private: // Converts a vector of replica groups to a vector of `XLA_FFI_ReplicaGroup`. // The results reference the id storage in `groups`, which must outlive them. diff --git a/third_party/xla/xla/ffi/api/collectives_c_api.h b/third_party/xla/xla/ffi/api/collectives_c_api.h index b635a92df6c44d..9a2e4719169d77 100644 --- a/third_party/xla/xla/ffi/api/collectives_c_api.h +++ b/third_party/xla/xla/ffi/api/collectives_c_api.h @@ -29,18 +29,22 @@ extern "C" { // XLA FFI Collectives API //===----------------------------------------------------------------------===// -// Exposes the XLA-owned host collective communicator to FFI handlers. +// Exposes the XLA-owned host collective communicator and collective memory +// window to FFI handlers. // `request_communicator` requests a clique in the Prepare stage; -// `get_communicator` returns the handle once cliques are acquired. The -// communicator is backend-defined and passed as an opaque `void*` (a handler on -// XLA:GPU reinterprets it as `ncclComm_t`, XLA:CPU as its own comm type, etc.). +// `get_communicator` returns the handle once cliques are acquired. +// `request_window` requests window registration for a batch of already- +// allocated buffers with the clique in Prepare; `get_window` returns the +// corresponding window handle once collective memory is acquired. +// The communicator and window are backend-defined and passed as opaque +// pointers (a handler on XLA:GPU reinterprets `XLA_FFI_Communicator*` as +// `ncclComm_t` and `XLA_FFI_Window*` as `ncclWindow_t`). // This API is backend agnostic; backend wiring lives in the runtime that // attaches the extension. -// Version macros for the collectives extension. #define XLA_FFI_Extension_Collectives 129 #define XLA_FFI_Extension_Collectives_MajorVersion 0 -#define XLA_FFI_Extension_Collectives_MinorVersion 1 +#define XLA_FFI_Extension_Collectives_MinorVersion 2 // Mirrors `xla::CollectiveOpGroupMode`. typedef enum XLA_FFI_CollectiveGroupMode { @@ -65,6 +69,10 @@ typedef struct XLA_FFI_CollectivesState XLA_FFI_CollectivesState; // type; callers reinterpret it (e.g. as `ncclComm_t`). typedef struct XLA_FFI_Communicator XLA_FFI_Communicator; +//===----------------------------------------------------------------------===// +// Communicator: request in Prepare, get in Initialize/Execute +//===----------------------------------------------------------------------===// + typedef struct XLA_FFI_Communicator_Request_Args { size_t struct_size; XLA_FFI_InternalExtension* extension_start; @@ -97,13 +105,81 @@ typedef struct XLA_FFI_Communicator_Get_Args { XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Communicator_Get_Args, communicator); -// Returns the non-owning communicator handle for the clique. Valid once cliques -// are acquired (Initialize/Execute stages). +// Returns the non-owning communicator handle for the clique. Valid once +// cliques are acquired (Initialize/Execute stages). typedef XLA_FFI_Error* XLA_FFI_Communicator_Get( const XLA_FFI_Collectives_Extension* self, XLA_FFI_Communicator_Get_Args* args); -// Main extension struct for the collectives API. +//===----------------------------------------------------------------------===// +// Collective memory window +//===----------------------------------------------------------------------===// +// +// Handlers request window registration for a batch of already-allocated +// buffers with the clique via `request_window` in Prepare, then look up an +// opaque `XLA_FFI_Window` per buffer via `get_window` in Init/Execute. +// Handlers reinterpret the window and call the backend's collective device APIs +// directly to obtain local, peer, and multicast pointers. + +// Opaque, non-owning collective memory window handle. The backend defines the +// concrete type; callers reinterpret it (e.g. as `ncclWindow_t`). +typedef struct XLA_FFI_Window XLA_FFI_Window; + +typedef struct XLA_FFI_CollectiveMemoryRegion { + const void* buffer; + size_t byte_size; + uint64_t flags; +} XLA_FFI_CollectiveMemoryRegion; + +typedef struct XLA_FFI_Window_Request_Args { + size_t struct_size; + XLA_FFI_InternalExtension* extension_start; + + XLA_FFI_CollectiveGroupMode group_mode; + const XLA_FFI_ReplicaGroup* groups; + size_t num_groups; + int64_t communication_id; + + const XLA_FFI_CollectiveMemoryRegion* regions; + size_t num_regions; +} XLA_FFI_Window_Request_Args; + +XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Window_Request_Args, num_regions); + +// Requests window registration for a batch of already-allocated buffers with +// the clique. Actual registration happens later once the clique is acquired; +// look up the resulting window handle with `get_window`. Prepare stage only. +typedef XLA_FFI_Error* XLA_FFI_Window_Request( + const XLA_FFI_Collectives_Extension* self, + XLA_FFI_Window_Request_Args* args); + +typedef struct XLA_FFI_Window_Get_Args { + size_t struct_size; + XLA_FFI_InternalExtension* extension_start; + + XLA_FFI_CollectiveGroupMode group_mode; + const XLA_FFI_ReplicaGroup* groups; + size_t num_groups; + int64_t communication_id; + + const void* buffer; // registered address (in) + XLA_FFI_Window* window; // out + size_t window_offset; // byte offset of `buffer` within `window` (out) +} XLA_FFI_Window_Get_Args; + +XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Window_Get_Args, window_offset); + +// Returns the non-owning collective memory window handle for a buffer whose +// registration was previously requested via `request_window`, along with the +// byte offset of `buffer` within that window. Valid once collective memory is +// acquired (Initialize/Execute stages). +typedef XLA_FFI_Error* XLA_FFI_Window_Get( + const XLA_FFI_Collectives_Extension* self, XLA_FFI_Window_Get_Args* args); + +//===----------------------------------------------------------------------===// +// Extension struct +//===----------------------------------------------------------------------===// + struct XLA_FFI_Collectives_Extension { XLA_FFI_Extension extension_base; @@ -111,9 +187,12 @@ struct XLA_FFI_Collectives_Extension { XLA_FFI_Communicator_Request* request_communicator; XLA_FFI_Communicator_Get* get_communicator; + + XLA_FFI_Window_Request* request_window; + XLA_FFI_Window_Get* get_window; }; -XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Collectives_Extension, get_communicator); +XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Collectives_Extension, get_window); #ifdef __cplusplus } // extern "C" diff --git a/third_party/xla/xla/ffi/api/collectives_ffi_test.cc b/third_party/xla/xla/ffi/api/collectives_ffi_test.cc index 397f24323a8372..95c487c54fd3a6 100644 --- a/third_party/xla/xla/ffi/api/collectives_ffi_test.cc +++ b/third_party/xla/xla/ffi/api/collectives_ffi_test.cc @@ -37,6 +37,9 @@ static const XLA_FFI_Api* Api() { return GetXlaFfiApi(); } static XLA_FFI_Communicator* const kFakeCommunicator = reinterpret_cast(0xC0FFEE); +static XLA_FFI_Window* const kFakeWindow = + reinterpret_cast(0xDEC0DE); +static constexpr size_t kFakeWindowOffset = 42; // Fake backend: implements the extension callbacks and records what the handler // requested so the test can check the C++ wrapper builds the C args correctly. @@ -47,6 +50,15 @@ struct FakeBackend { size_t requested_num_groups = 0; int64_t requested_communication_id = -1; std::vector requested_ids; + + bool request_window_called = false; + XLA_FFI_CollectiveGroupMode requested_window_group_mode = + XLA_FFI_GROUP_CROSS_REPLICA; + size_t requested_window_num_groups = 0; + int64_t requested_window_communication_id = -1; + std::vector requested_window_ids; + std::vector requested_regions; + const void* get_window_buffer = nullptr; }; static XLA_FFI_Error* FakeRequestCommunicator( @@ -73,18 +85,52 @@ static XLA_FFI_Error* FakeGetCommunicator( return nullptr; } +static XLA_FFI_Error* FakeRequestWindow( + const XLA_FFI_Collectives_Extension* self, + XLA_FFI_Window_Request_Args* args) { + auto* backend = reinterpret_cast(self->state); + backend->request_window_called = true; + backend->requested_window_group_mode = args->group_mode; + backend->requested_window_num_groups = args->num_groups; + backend->requested_window_communication_id = args->communication_id; + backend->requested_window_ids.clear(); + for (size_t i = 0; i < args->num_groups; ++i) { + for (size_t j = 0; j < args->groups[i].size; ++j) { + backend->requested_window_ids.push_back(args->groups[i].ids[j]); + } + } + backend->requested_regions.clear(); + for (size_t i = 0; i < args->num_regions; ++i) { + backend->requested_regions.push_back(args->regions[i]); + } + return nullptr; +} + +static XLA_FFI_Error* FakeGetWindow(const XLA_FFI_Collectives_Extension* self, + XLA_FFI_Window_Get_Args* args) { + auto* backend = reinterpret_cast(self->state); + backend->get_window_buffer = args->buffer; + args->window = kFakeWindow; + args->window_offset = kFakeWindowOffset; + return nullptr; +} + // Builds a collectives extension for the fake backend. Mirrors the small // builder each backend uses to publish the extension (see the GPU backend). static XLA_FFI_Collectives_Extension MakeFakeCollectivesExtension( XLA_FFI_CollectivesState* state, XLA_FFI_Communicator_Request* request_communicator, - XLA_FFI_Communicator_Get* get_communicator) { - XLA_FFI_Collectives_Extension ext; + XLA_FFI_Communicator_Get* get_communicator, + XLA_FFI_Window_Request* request_window = nullptr, + XLA_FFI_Window_Get* get_window = nullptr) { + XLA_FFI_Collectives_Extension ext = {}; ext.extension_base = MakeExtensionHeader>(); ext.state = state; ext.request_communicator = request_communicator; ext.get_communicator = get_communicator; + ext.request_window = request_window; + ext.get_window = get_window; return ext; } @@ -133,5 +179,66 @@ TEST(CollectivesFfiTest, RequestAndGetCommunicator) { EXPECT_THAT(backend.requested_ids, ElementsAre(0, 1)); } +TEST(CollectivesFfiTest, RequestAndGetWindow) { + static uint32_t fake_buffer0 = 0; + static uint32_t fake_buffer1 = 0; + const std::vector regions = { + {&fake_buffer0, sizeof(fake_buffer0)}, + {&fake_buffer1, sizeof(fake_buffer1)}, + }; + + bool called = false; + WindowLookup got_lookup = {nullptr, 0}; + + auto handler = Ffi::Bind().Ctx>().To( + [&](Communicator comm) -> Error { + called = true; + if (Error status = comm.RequestWindow(GroupMode::kFlattenedId, {{0, 1}}, + /*communication_id=*/9, regions); + status.failure()) { + return status; + } + ErrorOr lookup_or = + comm.GetWindow(GroupMode::kFlattenedId, {{0, 1}}, + /*communication_id=*/9, &fake_buffer0); + if (lookup_or.has_error()) { + return lookup_or.error(); + } + got_lookup = lookup_or.value(); + return Error::Success(); + }); + + CallFrameBuilder builder(/*num_args=*/0, /*num_rets=*/0); + auto call_frame = builder.Build(); + + FakeBackend backend; + XLA_FFI_Collectives_Extension ext = MakeFakeCollectivesExtension( + reinterpret_cast(&backend), + FakeRequestCommunicator, FakeGetCommunicator, FakeRequestWindow, + FakeGetWindow); + + InvokeContext context; + context.extension_start = &ext.extension_base; + + auto status = + Invoke(Api(), *handler, call_frame, context, ExecutionStage::kExecute); + + ASSERT_OK(status); + EXPECT_TRUE(called); + EXPECT_TRUE(backend.request_window_called); + EXPECT_EQ(backend.requested_window_group_mode, XLA_FFI_GROUP_FLATTENED_ID); + EXPECT_EQ(backend.requested_window_num_groups, 1u); + EXPECT_EQ(backend.requested_window_communication_id, 9); + EXPECT_THAT(backend.requested_window_ids, ElementsAre(0, 1)); + ASSERT_EQ(backend.requested_regions.size(), 2u); + EXPECT_EQ(backend.requested_regions[0].buffer, &fake_buffer0); + EXPECT_EQ(backend.requested_regions[0].byte_size, sizeof(fake_buffer0)); + EXPECT_EQ(backend.requested_regions[1].buffer, &fake_buffer1); + EXPECT_EQ(backend.requested_regions[1].byte_size, sizeof(fake_buffer1)); + EXPECT_EQ(backend.get_window_buffer, &fake_buffer0); + EXPECT_EQ(got_lookup.window, kFakeWindow); + EXPECT_EQ(got_lookup.offset, kFakeWindowOffset); +} + } // namespace } // namespace xla::ffi diff --git a/third_party/xla/xla/ffi/collectives_ffi_test.cc b/third_party/xla/xla/ffi/collectives_ffi_test.cc index 389b957cf9dbef..af29ebc81ffa23 100644 --- a/third_party/xla/xla/ffi/collectives_ffi_test.cc +++ b/third_party/xla/xla/ffi/collectives_ffi_test.cc @@ -39,6 +39,9 @@ static const XLA_FFI_Api* Api() { return GetXlaFfiApi(); } static XLA_FFI_Communicator* const kFakeCommunicator = reinterpret_cast(0xC0FFEE); +static XLA_FFI_Window* const kFakeWindow = + reinterpret_cast(0xDEC0DE); +static constexpr size_t kFakeWindowOffset = 42; // Fake backend: implements the extension callbacks and records what the handler // requested so the test can check the C++ wrapper builds the C args correctly. @@ -49,6 +52,15 @@ struct FakeBackend { size_t requested_num_groups = 0; int64_t requested_communication_id = -1; std::vector requested_ids; + + bool request_window_called = false; + XLA_FFI_CollectiveGroupMode requested_window_group_mode = + XLA_FFI_GROUP_CROSS_REPLICA; + size_t requested_window_num_groups = 0; + int64_t requested_window_communication_id = -1; + std::vector requested_window_ids; + std::vector requested_regions; + const void* get_window_buffer = nullptr; }; static XLA_FFI_Error* FakeRequestCommunicator( @@ -75,18 +87,52 @@ static XLA_FFI_Error* FakeGetCommunicator( return nullptr; } +static XLA_FFI_Error* FakeRequestWindow( + const XLA_FFI_Collectives_Extension* self, + XLA_FFI_Window_Request_Args* args) { + auto* backend = reinterpret_cast(self->state); + backend->request_window_called = true; + backend->requested_window_group_mode = args->group_mode; + backend->requested_window_num_groups = args->num_groups; + backend->requested_window_communication_id = args->communication_id; + backend->requested_window_ids.clear(); + for (size_t i = 0; i < args->num_groups; ++i) { + for (size_t j = 0; j < args->groups[i].size; ++j) { + backend->requested_window_ids.push_back(args->groups[i].ids[j]); + } + } + backend->requested_regions.clear(); + for (size_t i = 0; i < args->num_regions; ++i) { + backend->requested_regions.push_back(args->regions[i]); + } + return nullptr; +} + +static XLA_FFI_Error* FakeGetWindow(const XLA_FFI_Collectives_Extension* self, + XLA_FFI_Window_Get_Args* args) { + auto* backend = reinterpret_cast(self->state); + backend->get_window_buffer = args->buffer; + args->window = kFakeWindow; + args->window_offset = kFakeWindowOffset; + return nullptr; +} + // Builds a collectives extension for the fake backend. Mirrors the small // builder each backend uses to publish the extension (see the GPU backend). static XLA_FFI_Collectives_Extension MakeFakeCollectivesExtension( XLA_FFI_CollectivesState* state, XLA_FFI_Communicator_Request* request_communicator, - XLA_FFI_Communicator_Get* get_communicator) { - XLA_FFI_Collectives_Extension ext; + XLA_FFI_Communicator_Get* get_communicator, + XLA_FFI_Window_Request* request_window = nullptr, + XLA_FFI_Window_Get* get_window = nullptr) { + XLA_FFI_Collectives_Extension ext = {}; ext.extension_base = MakeExtensionHeader>(); ext.state = state; ext.request_communicator = request_communicator; ext.get_communicator = get_communicator; + ext.request_window = request_window; + ext.get_window = get_window; return ext; } @@ -135,5 +181,67 @@ TEST(CollectivesFfiTest, RequestAndGetCommunicator) { EXPECT_THAT(backend.requested_ids, ElementsAre(0, 1)); } +TEST(CollectivesFfiTest, RequestAndGetWindow) { + static uint32_t fake_buffer0 = 0; + static uint32_t fake_buffer1 = 0; + const std::vector regions = { + {&fake_buffer0, sizeof(fake_buffer0)}, + {&fake_buffer1, sizeof(fake_buffer1)}, + }; + + bool called = false; + WindowLookup got_lookup = {nullptr, 0}; + + auto handler = Ffi::Bind().Ctx>().To( + [&](Communicator comm) -> absl::Status { + called = true; + if (absl::Status status = + comm.RequestWindow(GroupMode::kFlattenedId, {{0, 1}}, + /*communication_id=*/9, regions); + !status.ok()) { + return status; + } + absl::StatusOr lookup_or = + comm.GetWindow(GroupMode::kFlattenedId, {{0, 1}}, + /*communication_id=*/9, &fake_buffer0); + if (!lookup_or.ok()) { + return lookup_or.status(); + } + got_lookup = *lookup_or; + return absl::OkStatus(); + }); + + CallFrameBuilder builder(/*num_args=*/0, /*num_rets=*/0); + auto call_frame = builder.Build(); + + FakeBackend backend; + XLA_FFI_Collectives_Extension ext = MakeFakeCollectivesExtension( + reinterpret_cast(&backend), + FakeRequestCommunicator, FakeGetCommunicator, FakeRequestWindow, + FakeGetWindow); + + InvokeContext context; + context.extension_start = &ext.extension_base; + + auto status = + Invoke(Api(), *handler, call_frame, context, ExecutionStage::kExecute); + + ASSERT_OK(status); + EXPECT_TRUE(called); + EXPECT_TRUE(backend.request_window_called); + EXPECT_EQ(backend.requested_window_group_mode, XLA_FFI_GROUP_FLATTENED_ID); + EXPECT_EQ(backend.requested_window_num_groups, 1u); + EXPECT_EQ(backend.requested_window_communication_id, 9); + EXPECT_THAT(backend.requested_window_ids, ElementsAre(0, 1)); + ASSERT_EQ(backend.requested_regions.size(), 2u); + EXPECT_EQ(backend.requested_regions[0].buffer, &fake_buffer0); + EXPECT_EQ(backend.requested_regions[0].byte_size, sizeof(fake_buffer0)); + EXPECT_EQ(backend.requested_regions[1].buffer, &fake_buffer1); + EXPECT_EQ(backend.requested_regions[1].byte_size, sizeof(fake_buffer1)); + EXPECT_EQ(backend.get_window_buffer, &fake_buffer0); + EXPECT_EQ(got_lookup.window, kFakeWindow); + EXPECT_EQ(got_lookup.offset, kFakeWindowOffset); +} + } // namespace } // namespace xla::ffi From 8d6ff1cda43d7280501c7b6446b2497eb9589b2b Mon Sep 17 00:00:00 2001 From: Shaogang Wang Date: Wed, 23 Sep 2026 06:21:49 -0700 Subject: [PATCH 06/28] PR #49155: [XLA:GPU] Preserve async ordering across command buffer boundaries MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/49155 ## πŸ“ Summary of Changes Command buffer conversion can split an async region and, in doing so, lose the ordering between operations queued on the same async stream. An async start/done pair runs its body on an async stream and joins it back to the main stream at the done. When a region cannot be captured (for example a collective that is disabled for command buffers), the pass still captured any convertible region nested inside it. That moves the nested body into a graph launched on the main stream, where nothing waits for the outer operation that is still running on the async stream. Notation: `CB{...}` is the set of thunks captured into one command buffer; minimum graph size is 1. **Example 1, nested pair on the same stream.** `A` and `B` both run on async stream `S`, so `A` must finish before `B` starts. `A` cannot be captured; `B` and the ordinary main-stream command `C` can. ```text Input: Start(A); Start(B); C; Done(B); Done(A) Before: Start(A); CB{Start(B); C; Done(B)}; Done(A) After: Start(A); Start(B); CB{C}; Done(B); Done(A) ``` Before, `B`'s body moves into the graph on the main stream, which has no wait for `A` to complete. After, `B` stays on stream `S` and `A β†’ B` is preserved, while `C` still benefits from capture. **Example 2, nested pair on a different stream (unchanged).** Same shape, but `B` runs on stream `T`. `B`'s body only ever waited on the main stream, never on `S`, so capturing it loses no ordering. The fix is scoped to the stream of the outstanding operation, so this case keeps the behaviour it has today. The same holds after an unmatched start such as a pipelined recv: only its own stream is affected. ```text Input: Start(A on S); Start(B on T); C; Done(B); Done(A) Before: Start(A); CB{Start(B); C; Done(B)}; Done(A) After: Start(A); CB{Start(B); C; Done(B)}; Done(A) ``` **Example 3, loop body joining an outer operation.** `A` cannot be captured; its done sits inside a loop body together with `C`, which reads `A`'s result. With `WHILE` enabled, the body alone looks convertible. ```text Input: Start(A); While{ Done(A); C } Before: Start(A); CB{ While{ Done(A); C } } After: Start(A); While{ Done(A); CB{C} } ``` Inside a graph the done is a no-op, so before this change `C` could run before `A` finishes. After, the done stays a thunk and only `C` is captured. This affects partially convertible workloads, such as compute overlapped with collectives that are disabled for command buffers. Part 1 of 3: #49158 moves LHS dependency construction into the emitter; #48900 then preserves async fork/join ordering inside captured command buffers. ## πŸš€ Kind of Contribution πŸ› Bug Fix ## πŸ§ͺ Unit Tests: `command_buffer_async_conversion_test` Copybara import of the project: -- ce5cde2dc8c619e58d6dcfd0220c8a523029c478 by Shawn Wang : [XLA:GPU] Preserve async ordering across command buffer boundaries -- deb8a3494242f599c95b91385e6ef4b4c56a03d1 by Shawn Wang : [XLA:GPU] Track open async streams during command buffer conversion Replace the positional open async region with a per-stream view. When a start stays a thunk, only its execution stream is open, and only until the done that joins it. Async regions on other streams can still be captured: their bodies only ever waited on the main stream, so a command buffer loses no ordering. A region is refused when it contains a nested start on an open stream, and loop or branch bodies are converted with the streams open at that point. Reject loop and branch bodies that contain an async done whose start is not in the same body. Capturing such a body whole would drop the join with the outstanding operation, while the top-level loop already keeps these thunks in place. -- 9986e62199ce6fb7662381bede9212bcbab09949 by Shawn Wang : [XLA:GPU] Fix include order in command buffer conversion pass clang-format places tsl/ headers after xla/ headers; the presubmit flagged the block touched by the previous commits. -- 53d4322cc4233c1e4c181bb0e6a1b29cbf03637b by Shawn Wang : [XLA:GPU] Tidy the command buffer async region handling Stop the nested-start walk at the first match instead of visiting the whole tree, inline the single-use region span helper, and let async start and done thunks fall through to the shared flush-and-append tail. Drop the unused status matcher include and dependency from the test. No behavior change. -- d9e5145b90e8f7d666ec4317359a57a978e1349e by Shawn Wang : [XLA:GPU] Check for duplicate outstanding async starts in conversion AsyncExecution::Start fails at run time when an execution is started again before its done, so a thunk sequence with two outstanding starts on one execution can never run. Make that explicit in AsyncRegionSize with a DCHECK instead of folding the case silently into the "no closed region" result; optimized builds keep leaving such a region uncaptured. Document the two remaining reasons the function returns zero. Reconcile the AsyncExecution::Start comment with its counter check: pipelined send/recv alternates Start and Done on a shared execution; it does not start it twice before a done. Merging this change closes #49155 PiperOrigin-RevId: 986711311 --- .../xla/xla/backends/gpu/runtime/BUILD | 29 ++ .../backends/gpu/runtime/async_execution.cc | 8 +- .../command_buffer_async_conversion_test.cc | 434 ++++++++++++++++++ .../runtime/command_buffer_conversion_pass.cc | 256 ++++++++--- .../runtime/command_buffer_conversion_pass.h | 12 + 5 files changed, 674 insertions(+), 65 deletions(-) create mode 100644 third_party/xla/xla/backends/gpu/runtime/command_buffer_async_conversion_test.cc diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD index 97f21128b1829e..bfe3b67572bec5 100644 --- a/third_party/xla/xla/backends/gpu/runtime/BUILD +++ b/third_party/xla/xla/backends/gpu/runtime/BUILD @@ -4779,6 +4779,7 @@ cc_library( srcs = ["command_buffer_conversion_pass.cc"], hdrs = ["command_buffer_conversion_pass.h"], deps = [ + ":async_execution", ":async_thunk", ":collective_group_thunk", ":command_buffer_cmd_emitter", @@ -4788,6 +4789,7 @@ cc_library( ":custom_call_thunk", ":device_to_device_copy_thunk", ":dynamic_slice_fusion_v2_thunk", + ":execution_stream_id", ":ragged_all_to_all_thunk", ":sequential_thunk", ":thunk", @@ -4802,8 +4804,10 @@ cc_library( "//xla/stream_executor:semantic_version", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:nullability", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/log:vlog_is_on", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", @@ -4817,6 +4821,31 @@ cc_library( ], ) +xla_cc_test( + name = "command_buffer_async_conversion_test", + srcs = ["command_buffer_async_conversion_test.cc"], + deps = [ + ":async_thunk", + ":command_buffer_conversion_pass", + ":command_buffer_thunk", + ":execution_stream_id", + ":memset_thunk", + ":replica_id_thunk", + ":thunk", + ":thunk_executor", + ":thunk_id", + ":thunk_pass_pipeline", + ":while_thunk", + "//xla:xla_proto_cc", + "//xla/service:buffer_assignment", + "//xla/service/gpu:gpu_device_info_for_tests", + "//xla/stream_executor:device_description", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_googletest//:gtest_main", + ], +) + xla_test( name = "command_buffer_conversion_pass_test", srcs = ["command_buffer_conversion_pass_test.cc"], diff --git a/third_party/xla/xla/backends/gpu/runtime/async_execution.cc b/third_party/xla/xla/backends/gpu/runtime/async_execution.cc index a5b45b78de9a32..26823b190a8800 100644 --- a/third_party/xla/xla/backends/gpu/runtime/async_execution.cc +++ b/third_party/xla/xla/backends/gpu/runtime/async_execution.cc @@ -138,9 +138,11 @@ absl::StatusOr AsyncExecution::Start( // Wait for all prior operations on `stream` before launching operations on // `async_stream`. We use a stream-level wait (not the shared event) so that // the event remains exclusively used for the asyncβ†’main completion signal. - // This is critical for pipelined send/recv where multiple Start() calls can - // happen before Done() (the event is safely overwritten on the async stream - // because the stream is ordered). + // This matters for pipelined send/recv, where several start thunks share one + // execution and Start()/Done() alternate across loop iterations: each Start() + // re-records the shared event on the ordered async stream after the previous + // Done() has waited on it. A second Start() before the matching Done() is + // rejected by the counter check above. ABSL_RETURN_IF_ERROR(async_stream->WaitFor(stream)); return ExecutionGuard(event, async_stream); diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_async_conversion_test.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_async_conversion_test.cc new file mode 100644 index 00000000000000..333dcaacb39701 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_async_conversion_test.cc @@ -0,0 +1,434 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include + +#include +#include +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "xla/backends/gpu/runtime/async_thunk.h" +#include "xla/backends/gpu/runtime/command_buffer_conversion_pass.h" +#include "xla/backends/gpu/runtime/command_buffer_thunk.h" +#include "xla/backends/gpu/runtime/execution_stream_id.h" +#include "xla/backends/gpu/runtime/memset_thunk.h" +#include "xla/backends/gpu/runtime/replica_id_thunk.h" +#include "xla/backends/gpu/runtime/thunk.h" +#include "xla/backends/gpu/runtime/thunk_executor.h" +#include "xla/backends/gpu/runtime/thunk_id.h" +#include "xla/backends/gpu/runtime/thunk_pass_pipeline.h" +#include "xla/backends/gpu/runtime/while_thunk.h" +#include "xla/service/buffer_assignment.h" +#include "xla/service/gpu/gpu_device_info_for_tests.h" +#include "xla/stream_executor/device_description.h" +#include "xla/xla.pb.h" + +namespace xla::gpu { +namespace { + +using ::testing::ElementsAre; +using ::testing::Pointee; + +MATCHER_P(ThunkKindIs, kind, "") { return arg.kind() == kind; } + +template +auto ThunkKindsAre(Kinds... kinds) { + return ElementsAre(Pointee(ThunkKindIs(kinds))...); +} + +class RejectingAllocator : public ThunkPassBufferAllocator { + public: + absl::StatusOr NewEmptyAllocation(int64_t size) override { + return absl::InternalError("Unexpected allocation during conversion"); + } +}; + +class CommandBufferAsyncConversionTest : public testing::Test { + protected: + Thunk::ThunkInfo Info() { + Thunk::ThunkInfo info; + info.thunk_id = ids_.GetNextThunkId(); + return info; + } + + void AddCommand(ThunkSequence& thunks) { + thunks.Emplace(Info(), slice_); + } + + AsyncStartThunk* Start(ThunkSequence& thunks, bool convertible = true, + ExecutionStreamId stream = ComputationStreamId(0)) { + ThunkSequence body; + if (convertible) { + AddCommand(body); + } else { + // Memset thunks are not eligible for conversion in this pass. + body.Emplace(Info(), 0, slice_); + } + auto start = + std::make_unique(Info(), stream, std::move(body)); + auto* result = start.get(); + thunks.push_back(std::move(start)); + return result; + } + + void Done(ThunkSequence& thunks, AsyncStartThunk* start) { + thunks.Emplace(Info(), start->async_execution()); + } + + WhileThunk* Loop(ThunkSequence& thunks, ThunkSequence body) { + auto loop = std::make_unique( + Info(), BufferAllocation::Slice(&allocation_, 0, 1), ThunkSequence{}, + std::move(body), /*trip_count=*/1); + auto* result = loop.get(); + thunks.push_back(std::move(loop)); + return result; + } + + const ThunkSequence& CommandBufferThunks(const Thunk& thunk) { + return static_cast(thunk).thunks()->thunks(); + } + + absl::StatusOr Convert(ThunkSequence& thunks, + bool enable_while = false) { + DebugOptions options; + options.add_xla_gpu_enable_command_buffer(DebugOptions::FUSION); + if (enable_while) { + options.add_xla_gpu_enable_command_buffer(DebugOptions::WHILE); + } + options.set_xla_gpu_graph_min_graph_size(1); + options.set_xla_gpu_command_buffer_scheduling_mode(DebugOptions::LHS); + se::DeviceDescription device = TestGpuDeviceInfo::RTXA6000DeviceInfo(); + RejectingAllocator allocator; + return CommandBufferConversionPass("test").Run( + &thunks, options, /*hlo_module=*/nullptr, device, allocator); + } + + BufferAllocation allocation_{0, sizeof(int32_t), 0}; + BufferAllocation::Slice slice_{&allocation_, 0, sizeof(int32_t)}; + ThunkIdGenerator ids_; +}; + +TEST_F(CommandBufferAsyncConversionTest, KeepsUnsupportedRegionsIntact) { + for (bool crossed : {false, true}) { + SCOPED_TRACE(crossed); + ThunkSequence thunks; + AddCommand(thunks); + auto* outer = Start(thunks, /*convertible=*/false); + auto* inner = Start(thunks); + Done(thunks, crossed ? outer : inner); + Done(thunks, crossed ? inner : outer); + AddCommand(thunks); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + // Capturing the inner start/done alone would drop its stream ordering + // against the unsupported outer operation. Capture resumes after both + // joins, and ordinary commands before the region can still be captured. + EXPECT_THAT(thunks, + ThunkKindsAre(Thunk::kCommandBuffer, Thunk::kAsyncStart, + Thunk::kAsyncStart, Thunk::kAsyncDone, + Thunk::kAsyncDone, Thunk::kCommandBuffer)); + } +} + +TEST_F(CommandBufferAsyncConversionTest, ConvertsPlainCommandsInOpenRegion) { + ThunkSequence thunks; + auto* start = Start(thunks, /*convertible=*/false); + AddCommand(thunks); + AddCommand(thunks); + Done(thunks, start); + AddCommand(thunks); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + // Commands between the unsupported start and its done execute on the main + // stream, exactly as their thunks would, so they can still be captured. + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kCommandBuffer, + Thunk::kAsyncDone, Thunk::kCommandBuffer)); + EXPECT_THAT(CommandBufferThunks(*thunks[1]), + ThunkKindsAre(Thunk::kReplicaId, Thunk::kReplicaId)); +} + +TEST_F(CommandBufferAsyncConversionTest, KeepsNestedRegionsInOpenRegion) { + ThunkSequence thunks; + auto* outer = Start(thunks, /*convertible=*/false); + auto* inner = Start(thunks); + AddCommand(thunks); + Done(thunks, inner); + AddCommand(thunks); + Done(thunks, outer); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + // Capturing the inner pair would run its body in a graph on the main stream + // and lose stream order with the outstanding outer operation. Its start and + // done stay thunks while the plain commands around them are captured. + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kAsyncStart, + Thunk::kCommandBuffer, Thunk::kAsyncDone, + Thunk::kCommandBuffer, Thunk::kAsyncDone)); +} + +TEST_F(CommandBufferAsyncConversionTest, ConvertsLoopBodiesInOpenRegion) { + for (bool body_has_async : {false, true}) { + SCOPED_TRACE(body_has_async); + ThunkSequence thunks; + auto* start = Start(thunks); + ThunkSequence body; + if (body_has_async) { + auto* inner = Start(body); + Done(body, inner); + } + AddCommand(body); + // WHILE conversion is disabled, so the loop keeps the region open and its + // body is converted recursively under the open-region rules. + WhileThunk* loop = Loop(thunks, std::move(body)); + Done(thunks, start); + AddCommand(thunks); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + EXPECT_THAT(thunks, + ThunkKindsAre(Thunk::kAsyncStart, Thunk::kWhile, + Thunk::kAsyncDone, Thunk::kCommandBuffer)); + if (body_has_async) { + EXPECT_THAT(loop->body_executor().thunks(), + ThunkKindsAre(Thunk::kAsyncStart, Thunk::kAsyncDone, + Thunk::kCommandBuffer)); + } else { + EXPECT_THAT(loop->body_executor().thunks(), + ThunkKindsAre(Thunk::kCommandBuffer)); + } + } +} + +TEST_F(CommandBufferAsyncConversionTest, + DoesNotCaptureLoopWithAsyncInOpenRegion) { + for (bool open_region : {false, true}) { + SCOPED_TRACE(open_region); + ThunkSequence thunks; + AsyncStartThunk* start = nullptr; + if (open_region) { + start = Start(thunks, /*convertible=*/false); + } + ThunkSequence body; + auto* inner = Start(body); + Done(body, inner); + WhileThunk* loop = Loop(thunks, std::move(body)); + if (open_region) { + Done(thunks, start); + } + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks, /*enable_while=*/true)); + if (open_region) { + // A convertible loop whose body starts async work on a stream with an + // outstanding operation must not become a graph on the main stream. + EXPECT_FALSE(changed); + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kWhile, + Thunk::kAsyncDone)); + EXPECT_THAT(loop->body_executor().thunks(), + ThunkKindsAre(Thunk::kAsyncStart, Thunk::kAsyncDone)); + } else { + EXPECT_TRUE(changed); + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kCommandBuffer)); + } + } +} + +TEST_F(CommandBufferAsyncConversionTest, ConvertsTailAfterUnmatchedStart) { + ThunkSequence thunks; + AddCommand(thunks); + Start(thunks); // Its done can belong to a different pipelined computation. + AddCommand(thunks); + auto* inner = Start(thunks); + Done(thunks, inner); + AddCommand(thunks); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + // The stream stays open until the end of the sequence: later pairs on the + // same stream remain thunks while plain commands are still captured. + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kCommandBuffer, Thunk::kAsyncStart, + Thunk::kCommandBuffer, Thunk::kAsyncStart, + Thunk::kAsyncDone, Thunk::kCommandBuffer)); +} + +TEST_F(CommandBufferAsyncConversionTest, + ConvertsTailAfterUnmatchedStartOnOtherStream) { + ThunkSequence thunks; + AddCommand(thunks); + Start(thunks, /*convertible=*/true, ComputationStreamId(0)); + AddCommand(thunks); + auto* inner = Start(thunks, /*convertible=*/true, ComputationStreamId(1)); + Done(thunks, inner); + AddCommand(thunks); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + // The unmatched start only keeps its own stream open. A pair on another + // stream never observed that stream in thunk mode either, so capturing it + // loses no ordering. + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kCommandBuffer, Thunk::kAsyncStart, + Thunk::kCommandBuffer)); + EXPECT_THAT(CommandBufferThunks(*thunks[2]), + ThunkKindsAre(Thunk::kReplicaId, Thunk::kAsyncStart, + Thunk::kAsyncDone, Thunk::kReplicaId)); +} + +TEST_F(CommandBufferAsyncConversionTest, ConvertsNestedRegionOnOtherStream) { + ThunkSequence thunks; + auto* outer = Start(thunks, /*convertible=*/false, ComputationStreamId(0)); + auto* inner = Start(thunks, /*convertible=*/true, ComputationStreamId(1)); + AddCommand(thunks); + Done(thunks, inner); + AddCommand(thunks); + Done(thunks, outer); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kCommandBuffer, + Thunk::kAsyncDone)); + EXPECT_THAT(CommandBufferThunks(*thunks[1]), + ThunkKindsAre(Thunk::kAsyncStart, Thunk::kReplicaId, + Thunk::kAsyncDone, Thunk::kReplicaId)); +} + +TEST_F(CommandBufferAsyncConversionTest, ClosesStreamAtMatchingDone) { + ThunkSequence thunks; + auto* a = Start(thunks, /*convertible=*/false, ComputationStreamId(0)); + auto* b = Start(thunks, /*convertible=*/false, ComputationStreamId(1)); + Done(thunks, a); + auto* c = Start(thunks, /*convertible=*/true, ComputationStreamId(0)); + Done(thunks, c); + Done(thunks, b); + AddCommand(thunks); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + // Stream 0 is joined at `a`'s done, so `c` can be captured even though + // stream 1 is still open. + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kAsyncStart, + Thunk::kAsyncDone, Thunk::kCommandBuffer, + Thunk::kAsyncDone, Thunk::kCommandBuffer)); + EXPECT_THAT(CommandBufferThunks(*thunks[3]), + ThunkKindsAre(Thunk::kAsyncStart, Thunk::kAsyncDone)); +} + +TEST_F(CommandBufferAsyncConversionTest, KeepsRegionWithStartOnOpenStream) { + ThunkSequence thunks; + auto* a = Start(thunks, /*convertible=*/false, ComputationStreamId(0)); + auto* b = Start(thunks, /*convertible=*/true, ComputationStreamId(1)); + auto* c = Start(thunks, /*convertible=*/true, ComputationStreamId(0)); + Done(thunks, c); + Done(thunks, b); + Done(thunks, a); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + // `b`'s region is closed and convertible, but it contains `c`, which starts + // work on the stream still occupied by `a`. Capturing the region would run + // `c`'s body in a graph on the main stream. + EXPECT_FALSE(changed); + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kAsyncStart, + Thunk::kAsyncStart, Thunk::kAsyncDone, + Thunk::kAsyncDone, Thunk::kAsyncDone)); +} + +TEST_F(CommandBufferAsyncConversionTest, ConvertsLoopBodyRegionOnOtherStream) { + ThunkSequence thunks; + auto* start = Start(thunks, /*convertible=*/false, ComputationStreamId(0)); + ThunkSequence body; + auto* inner = Start(body, /*convertible=*/true, ComputationStreamId(1)); + Done(body, inner); + AddCommand(body); + WhileThunk* loop = Loop(thunks, std::move(body)); + Done(thunks, start); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kWhile, + Thunk::kAsyncDone)); + // Only stream 0 is open inside the loop body; the pair on stream 1 is + // captured together with the plain command. + EXPECT_THAT(loop->body_executor().thunks(), + ThunkKindsAre(Thunk::kCommandBuffer)); + EXPECT_THAT( + CommandBufferThunks(*loop->body_executor().thunks()[0]), + ThunkKindsAre(Thunk::kAsyncStart, Thunk::kAsyncDone, Thunk::kReplicaId)); +} + +TEST_F(CommandBufferAsyncConversionTest, DoesNotCaptureLoopWithLoneDone) { + ThunkSequence thunks; + auto* start = Start(thunks, /*convertible=*/false); + ThunkSequence body; + Done(body, start); + AddCommand(body); + WhileThunk* loop = Loop(thunks, std::move(body)); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks, /*enable_while=*/true)); + EXPECT_TRUE(changed); + // The loop body joins an operation started outside the loop. Capturing the + // loop whole would drop that join, so only the plain command in the body is + // captured. + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kWhile)); + EXPECT_THAT(loop->body_executor().thunks(), + ThunkKindsAre(Thunk::kAsyncDone, Thunk::kCommandBuffer)); +} + +TEST_F(CommandBufferAsyncConversionTest, MatchesCanonicalAsyncExecution) { + ThunkSequence thunks; + auto* canonical = Start(thunks); + Done(thunks, canonical); + ThunkSequence body; + AddCommand(body); + thunks.push_back(std::make_unique( + Info(), ComputationStreamId(0), std::move(body), + canonical->async_execution())); + Done(thunks, canonical); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kCommandBuffer)); +} + +TEST_F(CommandBufferAsyncConversionTest, DuplicateOutstandingStart) { + ThunkSequence thunks; + auto* canonical = Start(thunks); + ThunkSequence body; + AddCommand(body); + thunks.push_back(std::make_unique( + Info(), ComputationStreamId(0), std::move(body), + canonical->async_execution())); + Done(thunks, canonical); + Done(thunks, canonical); + AddCommand(thunks); + + // The runtime permits only one outstanding start per AsyncExecution, so this + // sequence can never execute. Debug builds catch it in the pass; optimized + // builds leave the region as thunks and still capture the trailing command. +#ifndef NDEBUG + EXPECT_DEATH(Convert(thunks).IgnoreError(), "started twice"); +#else + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kAsyncStart, + Thunk::kAsyncDone, Thunk::kAsyncDone, + Thunk::kCommandBuffer)); +#endif +} + +} // namespace +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc index 3730af6a1cd1db..295772501d1968 100644 --- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc +++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc @@ -29,7 +29,9 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/base/nullability.h" +#include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" +#include "absl/log/check.h" #include "absl/log/log.h" #include "absl/log/vlog_is_on.h" #include "absl/status/status.h" @@ -38,6 +40,7 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "absl/types/span.h" +#include "xla/backends/gpu/runtime/async_execution.h" #include "xla/backends/gpu/runtime/async_thunk.h" #include "xla/backends/gpu/runtime/collective_group_thunk.h" #include "xla/backends/gpu/runtime/command_buffer_cmd_emitter.h" @@ -47,6 +50,7 @@ limitations under the License. #include "xla/backends/gpu/runtime/custom_call_thunk.h" #include "xla/backends/gpu/runtime/device_to_device_copy_thunk.h" #include "xla/backends/gpu/runtime/dynamic_slice_fusion_v2_thunk.h" +#include "xla/backends/gpu/runtime/execution_stream_id.h" #include "xla/backends/gpu/runtime/ragged_all_to_all_thunk.h" #include "xla/backends/gpu/runtime/sequential_thunk.h" #include "xla/backends/gpu/runtime/thunk.h" @@ -459,19 +463,23 @@ bool ThunkSequenceIsConvertible(const ThunkSequence& thunks, return false; } i += region_size - 1; + } else if (thunk->kind() == Thunk::kAsyncDone) { + // Every done that belongs to a start in this sequence was consumed above + // as part of its region. This done joins an operation started outside + // the sequence; capturing it would drop the join (the top-level loop in + // `RunImpl` keeps such thunks in place for the same reason). + return false; } } return true; } // Collects and returns the size of the shortest non-empty sequence of thunks -// that form a valid async region. -// The sequence considered as a valid async region if each start thunk has a -// corresponding done thunk and vice versa, and all thunks in between are -// convertible. If there is another start thunk between the original start and -// done, we may potentially extend the sequence to include its corresponding -// done thunk. For example, if we call this function on async-start_a in the -// following sequence: +// that form a closed async region: each start thunk has a corresponding done +// thunk and vice versa. If there is another start thunk between the original +// start and done, we may potentially extend the sequence to include its +// corresponding done thunk. For example, if we call this function on +// async-start_a in the following sequence: // // async_start_a // async_start_b @@ -480,48 +488,113 @@ bool ThunkSequenceIsConvertible(const ThunkSequence& thunks, // // The returned sequence will contain async_done_b. So that all async pairs // are captured by the same command buffer. -size_t CheckAsyncRegion(absl::Span> thunks, - const CommandBufferConfig& config) { - absl::flat_hash_set unpaired_ids; +// Find the boundary independently of command-buffer eligibility. If any thunk +// is unsupported, the whole region must remain outside the command buffer: +// capturing an inner region could lose ordering with an outstanding operation +// on the same async stream. +// +// Returns 0 when no closed region starts here: either a start in the scanned +// range has no done in this sequence, or a done in the range joins a start +// from outside it. The caller handles both the same way, by leaving the start +// as a thunk and keeping its stream open until it is joined. +size_t AsyncRegionSize(absl::Span> thunks) { + absl::flat_hash_set unpaired_executions; for (size_t i = 0; i < thunks.size(); ++i) { auto& thunk = thunks[i]; - // Check if thunk is convertible - if (!IsConvertible(*thunk, config)) { - return 0; // All thunks in the region must be convertible. - } - - // Track AsyncStartThunk/AsyncDoneThunk pairs via AsyncExecutionId. if (thunk->kind() == Thunk::kAsyncStart) { - unpaired_ids.insert(static_cast(*thunk) - .async_execution_id() - .value()); + // Pipelined starts can share the canonical start's execution state, but + // AsyncExecution::Start rejects a second start before the matching done, + // so a valid sequence never inserts the same execution twice. Optimized + // builds fall back to leaving such a region uncaptured. + bool inserted = unpaired_executions + .insert(static_cast(*thunk) + .async_execution() + .get()) + .second; + DCHECK(inserted) << "Async execution started twice before its done: " + << thunk->profile_annotation(); + if (!inserted) { + return 0; + } } if (thunk->kind() == Thunk::kAsyncDone) { - auto id = static_cast(*thunk) - .async_execution_id() - .value(); - auto it = unpaired_ids.find(id); - if (it == unpaired_ids.end()) { + auto* execution = + static_cast(*thunk).async_execution().get(); + auto it = unpaired_executions.find(execution); + if (it == unpaired_executions.end()) { return 0; // Done without matching start in the region. } - unpaired_ids.erase(it); + unpaired_executions.erase(it); } - if (unpaired_ids.empty()) { + if (unpaired_executions.empty()) { return i + 1; // All start/done pairs are matched. } } - return 0; // error didn't find an end for some start + return 0; // At least one start has no matching done in this sequence. } -// Returns the shortest non-empty sequence of thunks that form a valid async -// region as a span. If no such region is found, an empty span is returned. -absl::Span> CollectAndCheckAsyncRegion( - absl::Span> thunks, - const CommandBufferConfig& config) { - return thunks.subspan(0, CheckAsyncRegion(thunks, config)); +// Returns the size of a closed region only if every thunk can be converted. +size_t CheckAsyncRegion(absl::Span> thunks, + const CommandBufferConfig& config) { + size_t size = AsyncRegionSize(thunks); + for (const std::unique_ptr& thunk : thunks.first(size)) { + if (!IsConvertible(*thunk, config)) { + return 0; + } + } + return size; +} + +// Returns true if `thunk` or any thunk nested in it starts an async region on +// one of `streams`. DynamicSliceFusionV2Thunk hides its embedded thunks from +// Thunk::Walk, so they are inspected explicitly. +bool ContainsAsyncStartOnStreams( + const Thunk& thunk, const absl::flat_hash_set& streams) { + if (streams.empty()) { + return false; + } + // A non-OK status stops the walk at the first match. + absl::Status walk = thunk.Walk([&](const Thunk* nested) -> absl::Status { + if (nested->kind() == Thunk::kAsyncStart) { + if (streams.contains(static_cast(*nested) + .execution_stream_id())) { + return absl::CancelledError(); + } + } else if (nested->kind() == Thunk::kDynamicSliceFusion) { + const auto& fusion = + static_cast(*nested); + for (const std::unique_ptr& embedded : fusion.thunks()) { + if (ContainsAsyncStartOnStreams(*embedded, streams)) { + return absl::CancelledError(); + } + } + } + return absl::OkStatus(); + }); + return !walk.ok(); +} + +// Returns the index one past the done that joins the start at +// `thunks[start_index]`, or `thunks.size()` if no thunk in this sequence joins +// it. Until that index, work started on the start's stream is outstanding. +size_t AsyncJoinEnd(absl::Span> thunks, + size_t start_index) { + const AsyncExecution* execution = + static_cast(*thunks[start_index]) + .async_execution() + .get(); + for (size_t i = start_index + 1; i < thunks.size(); ++i) { + if (thunks[i]->kind() == Thunk::kAsyncDone && + static_cast(*thunks[i]) + .async_execution() + .get() == execution) { + return i + 1; + } + } + return thunks.size(); } absl::StatusOr GetSynchronizationMode( @@ -671,6 +744,16 @@ absl::StatusOr CommandBufferConversionPass::Run( const HloModule* absl_nullable hlo_module, const se::DeviceDescription& device_info, ThunkPassBufferAllocator& allocator) { + return RunImpl(thunk_sequence, debug_options, hlo_module, device_info, + allocator, /*open_async_streams=*/{}); +} + +absl::StatusOr CommandBufferConversionPass::RunImpl( + ThunkSequence* thunk_sequence, const DebugOptions& debug_options, + const HloModule* absl_nullable hlo_module, + const se::DeviceDescription& device_info, + ThunkPassBufferAllocator& allocator, + const absl::flat_hash_set& open_async_streams) { tsl::profiler::TraceMe traceme("CommandBufferConversionPass"); CommandBufferConfig config = @@ -694,54 +777,103 @@ absl::StatusOr CommandBufferConversionPass::Run( auto& original_thunks = *thunk_sequence; + // An async region is "open" when it cannot be captured as a whole: some + // thunk in it is not convertible, or its start has no matching done in this + // sequence. Its start stays a thunk and enqueues work on its async stream + // that a command buffer launched on the main stream cannot observe. That + // stream is "open" until the done that joins it, so while it is open: + // * plain commands still execute on the main stream and can be captured, + // exactly as their thunks never observed the async stream either; + // * async regions on other streams can still be captured: their bodies + // only ever waited on the main stream, so a graph loses no ordering; + // * async starts on the open stream stay in place, because capturing them + // would move their bodies into a graph on the main stream and lose + // stream order with the outstanding operation; + // * control flow containing such starts is not captured whole; its bodies + // are converted recursively with the same streams open throughout. + // `open_stream_ends` maps each open stream to the index one past the thunk + // that joins it, or the end of the sequence for an unmatched start. + absl::flat_hash_map open_stream_ends; + for (ExecutionStreamId stream : open_async_streams) { + open_stream_ends[stream] = original_thunks.size(); + } + auto open_streams_at = [&](size_t index) { + absl::flat_hash_set streams; + for (const auto& [stream, end] : open_stream_ends) { + if (index < end) { + streams.insert(stream); + } + } + return streams; + }; + for (size_t i = 0; i < original_thunks.size(); ++i) { auto& thunk = original_thunks[i]; + const absl::flat_hash_set open_streams = + open_streams_at(i); - // We always have to capture both corresponding start and done events in the - // same command buffer. if (thunk->kind() == Thunk::kAsyncStart) { - // Collect and check async region - absl::Span> region = CollectAndCheckAsyncRegion( - absl::MakeSpan(original_thunks).subspan(i), config); - - if (!region.empty()) { - // If a valid region is found, add the whole region to the current - // sequence and continue processing. - i += region.size() - 1; - absl::c_move(region, std::back_inserter(current_command_buffer_thunks)); - continue; + const auto& start = static_cast(*thunk); + // We always have to capture both corresponding start and done events in + // the same command buffer. + if (!open_streams.contains(start.execution_stream_id())) { + absl::Span> tail = + absl::MakeSpan(original_thunks).subspan(i); + absl::Span> region = + tail.first(CheckAsyncRegion(tail, config)); + if (!region.empty() && + absl::c_none_of(region, [&](const std::unique_ptr& nested) { + return ContainsAsyncStartOnStreams(*nested, open_streams); + })) { + // If a valid region is found, add the whole region to the current + // sequence and continue processing. + i += region.size() - 1; + absl::c_move(region, + std::back_inserter(current_command_buffer_thunks)); + continue; + } } - } else if (IsConvertible(*thunk.get(), config) && - thunk->kind() != Thunk::kAsyncDone) { - // Check if thunk is convertible and not an async done: async done thunks - // can be only added to the current_command_buffer_thunks as part of a - // valid async regions. + // This start stays a thunk so that its body keeps executing on its async + // stream; the stream is open until the done that joins it. + size_t& end = open_stream_ends[start.execution_stream_id()]; + end = std::max(end, AsyncJoinEnd(original_thunks, i)); + } + + // Async start and done thunks are only captured as part of a valid async + // region above; on their own they stay in place. + const bool is_async_boundary = thunk->kind() == Thunk::kAsyncStart || + thunk->kind() == Thunk::kAsyncDone; + if (!is_async_boundary && IsConvertible(*thunk, config) && + !ContainsAsyncStartOnStreams(*thunk, open_streams)) { current_command_buffer_thunks.push_back(std::move(thunk)); continue; } + if (thunk->kind() == Thunk::kWhile) { - // If a `WhileThunk` itself is not eligible for conversion into a - // command buffer, we attempt to convert thunks within its body + // If a `WhileThunk` itself is not captured into a command buffer, we + // attempt to convert thunks within its body. auto while_thunk = static_cast(thunk.get()); - ABSL_ASSIGN_OR_RETURN(bool changed_in_body, - Run(&while_thunk->body_executor().thunks(), - debug_options, hlo_module, device_info, allocator)); + ABSL_ASSIGN_OR_RETURN( + bool changed_in_body, + RunImpl(&while_thunk->body_executor().thunks(), debug_options, + hlo_module, device_info, allocator, open_streams)); changed |= changed_in_body; } else if (thunk->kind() == Thunk::kConditional) { - // If a `ConditionalThunk` itself is not eligible for conversion into a - // command buffer, we attempt to convert thunks within its branches. + // If a `ConditionalThunk` itself is not captured into a command buffer, + // we attempt to convert thunks within its branches. auto conditional_thunk = static_cast(thunk.get()); for (auto& branch_executor : conditional_thunk->branch_executors()) { - ABSL_ASSIGN_OR_RETURN(bool changed_in_branch, - Run(&branch_executor.thunks(), debug_options, - hlo_module, device_info, allocator)); + ABSL_ASSIGN_OR_RETURN( + bool changed_in_branch, + RunImpl(&branch_executor.thunks(), debug_options, hlo_module, + device_info, allocator, open_streams)); changed |= changed_in_branch; } } - // If the current thunk is not convertible, flush collected eligible thunk + // If the current thunk is not captured, flush collected eligible thunks // to a command buffer thunk and add it to the processed sequence. Then add - // non-convertible thunk to the sequence. + // the thunk itself to the sequence. ABSL_RETURN_IF_ERROR(flush_command_buffer()); new_thunks.push_back(std::move(thunk)); } diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.h b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.h index c7cff299b8cb16..30e74f21fd2885 100644 --- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.h +++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.h @@ -24,6 +24,7 @@ limitations under the License. #include "absl/container/flat_hash_set.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" +#include "xla/backends/gpu/runtime/execution_stream_id.h" #include "xla/backends/gpu/runtime/thunk.h" #include "xla/backends/gpu/runtime/thunk_pass_pipeline.h" #include "xla/hlo/ir/hlo_module.h" @@ -66,6 +67,17 @@ class CommandBufferConversionPass : public ThunkPassInterface { }; private: + // Implementation of `Run`. `open_async_streams` are the async streams with + // an outstanding operation for the whole of `thunk_sequence`: it is the body + // of a loop or branch nested in an async region that is not captured as a + // whole; see the comment in `RunImpl`. + absl::StatusOr RunImpl( + ThunkSequence* thunk_sequence, const DebugOptions& debug_options, + const HloModule* absl_nullable hlo_module, + const se::DeviceDescription& device_info, + ThunkPassBufferAllocator& allocator, + const absl::flat_hash_set& open_async_streams); + std::string module_name_; }; From ff5040cdee9e07b60417b9d3a0a01e3aa18227e0 Mon Sep 17 00:00:00 2001 From: Henning Becker Date: Wed, 23 Sep 2026 06:24:21 -0700 Subject: [PATCH 07/28] Document the native custom-call handler API. The thunk-folding API now has enough moving parts that reading the headers is no longer a reasonable way to learn it: a handler author has to know when to prefer a native handler over XLA FFI, which of the three index spaces around a kernel launch a given number lives in, why buffers must be resolved through the emitter context rather than through the buffer assignment, and how to test the result without a GPU. Add a README covering all of that, with the `write_value_thunk_folded` demonstrator as a worked example, and state the two things the API deliberately does not do: validate operand element types, and let a handler decline an instruction and fall back to the FFI path. Also drop the pointer to `g3doc/native_custom_call_handlers.md` from `native_custom_call_handler_utils.h`. That path never existed. PiperOrigin-RevId: 986712400 --- .../native_custom_call_handler_utils.h | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_handler_utils.h b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_handler_utils.h index ea80082b1c0447..37bf1330934fee 100644 --- a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_handler_utils.h +++ b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_handler_utils.h @@ -34,8 +34,7 @@ limitations under the License. #include "xla/stream_executor/launch_dim.h" // Helpers for writing custom-call thunk-folding handlers. See -// `native_custom_call_handler_registry.h` for the mechanism itself and -// `g3doc/native_custom_call_handlers.md` for a tutorial. +// `native_custom_call_handler_registry.h` for the mechanism itself. namespace xla::gpu { From b48bf7cbfe2df27456a4955c589e72a535f84fba Mon Sep 17 00:00:00 2001 From: Bhavani Subramanian Date: Wed, 23 Sep 2026 06:27:59 -0700 Subject: [PATCH 08/28] PR #49199: [XLA:GPU][oneAPI] Fix unit test failure due to platform version Imported from GitHub PR https://github.com/openxla/xla/pull/49199 This PR fixes `StreamExecutorGpuClientTest.PlatformVersionIsDerivedAtRuntime` for the oneAPI backend by adding `oneapi` to the list of accepted platform-version prefixes. It also fixes a crash when running all tests in `StreamExecutorGpuClientTest`, caused by an uncaught SYCL exception during `SyclStream` teardown. Copybara import of the project: -- 231dd5a0e5a3a6b10d2a39def262a1ae0c018aec by Bhavani Subramanian : Fix unit test failure due to platform version on oneAPI backend Merging this change closes #49199 PiperOrigin-RevId: 986714208 --- .../xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc | 3 ++- .../xla/xla/stream_executor/sycl/sycl_stream.cc | 14 ++++++++++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc index bd2f82bf978755..2feae674faa7b3 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc @@ -344,7 +344,8 @@ TEST(StreamExecutorGpuClientTest, PlatformVersionIsDerivedAtRuntime) { const absl::string_view version = client->platform_version(); EXPECT_NE(version, ""); EXPECT_TRUE(absl::StartsWith(version, "cuda ") || - absl::StartsWith(version, "rocm ")) + absl::StartsWith(version, "rocm ") || + absl::StartsWith(version, "oneapi")) << "unexpected platform version: " << version; } diff --git a/third_party/xla/xla/stream_executor/sycl/sycl_stream.cc b/third_party/xla/xla/stream_executor/sycl/sycl_stream.cc index f1c1220d911487..64195ed43ed918 100644 --- a/third_party/xla/xla/stream_executor/sycl/sycl_stream.cc +++ b/third_party/xla/xla/stream_executor/sycl/sycl_stream.cc @@ -234,7 +234,13 @@ absl::Status SyclStream::DoHostCallbackWithStatus( } absl::Status SyclStream::BlockHostUntilDone() { - stream_handle_->wait(); + // Called from ~SyclStream(), so exceptions must be caught, not propagated. + try { + stream_handle_->wait(); + } catch (const ::sycl::exception& e) { + return absl::InternalError(absl::StrCat( + "SYCL exception in SyclStream::BlockHostUntilDone: ", e.what())); + } return absl::OkStatus(); } @@ -272,7 +278,11 @@ absl::StatusOr> SyclStream::Create( SyclStream::~SyclStream() { // Wait for all pending operations to complete before destroying the stream. - BlockHostUntilDone().IgnoreError(); + absl::Status wait_status = BlockHostUntilDone(); + if (!wait_status.ok()) { + LOG(ERROR) << "BlockHostUntilDone failed during stream destruction: " + << wait_status; + } // Remove this stream from the executor's list of allocated streams. executor_->DeallocateStream(this); From 294522f7acc83f81dbfb8eeae26f47e3028e0dad Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 23 Sep 2026 06:33:07 -0700 Subject: [PATCH 09/28] Automated Code Change PiperOrigin-RevId: 986716923 --- third_party/xla/xla/hlo/transforms/collectives/BUILD | 8 ++++++++ .../hlo/transforms/collectives/all_gather_decomposer.h | 1 + .../transforms/collectives/async_collective_replacer.cc | 1 + .../collectives/collective_permute_decomposer.cc | 2 ++ .../collectives/collective_permute_decomposer.h | 1 + .../collectives/collective_permute_decomposer_test.cc | 2 ++ .../collectives/reduce_scatter_combiner_test.cc | 1 + .../transforms/collectives/reduce_scatter_decomposer.cc | 1 + .../collectives/reduce_scatter_decomposer_test.cc | 1 + 9 files changed, 18 insertions(+) diff --git a/third_party/xla/xla/hlo/transforms/collectives/BUILD b/third_party/xla/xla/hlo/transforms/collectives/BUILD index d80191f24c2867..51a74a2a0043dd 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/BUILD +++ b/third_party/xla/xla/hlo/transforms/collectives/BUILD @@ -102,6 +102,7 @@ cc_library( "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/log:vlog_is_on", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings:string_view", ], @@ -785,6 +786,8 @@ cc_library( deps = [ ":collective_permute_cycle", "//xla:shape_util", + "//xla:xla_data_proto_cc", + "//xla:xla_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", "//xla/service:call_graph", @@ -808,6 +811,8 @@ xla_cc_test( srcs = ["collective_permute_decomposer_test.cc"], deps = [ ":collective_permute_decomposer", + "//xla:xla_data_proto_cc", + "//xla:xla_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/service:collective_ops_utils", @@ -865,6 +870,7 @@ xla_cc_test( srcs = ["reduce_scatter_combiner_test.cc"], deps = [ ":reduce_scatter_combiner", + "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/hlo/utils:hlo_matchers", @@ -961,6 +967,7 @@ cc_library( hdrs = ["reduce_scatter_decomposer.h"], deps = [ "//xla:shape_util", + "//xla:xla_data_proto_cc", "//xla/hlo/ir:collective_op_group_mode", "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", @@ -982,6 +989,7 @@ xla_cc_test( "//xla:literal", "//xla:literal_util", "//xla:shape_util", + "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/hlo/utils:hlo_matchers", diff --git a/third_party/xla/xla/hlo/transforms/collectives/all_gather_decomposer.h b/third_party/xla/xla/hlo/transforms/collectives/all_gather_decomposer.h index 84ef994189c11e..55ac3c62864f1a 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/all_gather_decomposer.h +++ b/third_party/xla/xla/hlo/transforms/collectives/all_gather_decomposer.h @@ -30,6 +30,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/pass/hlo_pass_interface.h" #include "xla/shape.h" +#include "xla/xla_data.pb.h" namespace xla { diff --git a/third_party/xla/xla/hlo/transforms/collectives/async_collective_replacer.cc b/third_party/xla/xla/hlo/transforms/collectives/async_collective_replacer.cc index 8433b2ee83a818..019c5a8cceca1a 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/async_collective_replacer.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/async_collective_replacer.cc @@ -23,6 +23,7 @@ limitations under the License. #include "absl/log/check.h" #include "absl/log/log.h" #include "absl/log/vlog_is_on.h" +#include "absl/status/status_macros.h" #include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" diff --git a/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer.cc b/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer.cc index 780927fbb9e9c9..83991e68a13546 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer.cc @@ -45,6 +45,8 @@ limitations under the License. #include "xla/service/source_target_pairs.h" #include "xla/shape.h" #include "xla/shape_util.h" +#include "xla/xla.pb.h" +#include "xla/xla_data.pb.h" namespace xla { diff --git a/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer.h b/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer.h index 1e963f2b451133..ce05ce5a5a936b 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer.h +++ b/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer.h @@ -23,6 +23,7 @@ limitations under the License. #include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/pass/hlo_pass_interface.h" +#include "xla/xla.pb.h" namespace xla { diff --git a/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer_test.cc b/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer_test.cc index c924facb525c83..40d713710554b2 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer_test.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer_test.cc @@ -32,6 +32,8 @@ limitations under the License. #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/service/collective_ops_utils.h" +#include "xla/xla.pb.h" +#include "xla/xla_data.pb.h" namespace xla { namespace { diff --git a/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_combiner_test.cc b/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_combiner_test.cc index 20f5a515662506..1f5c7ef3666628 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_combiner_test.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_combiner_test.cc @@ -32,6 +32,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/utils/hlo_matchers.h" +#include "xla/xla_data.pb.h" namespace xla { namespace { diff --git a/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_decomposer.cc b/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_decomposer.cc index 09cdb0924c9e2e..57c6c2f20a3539 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_decomposer.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_decomposer.cc @@ -34,6 +34,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/utils/hlo_query.h" #include "xla/service/collective_decomposer_utils.h" +#include "xla/xla_data.pb.h" namespace xla { diff --git a/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_decomposer_test.cc b/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_decomposer_test.cc index 41e6bdbf77dab4..0e6dc41f4a2deb 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_decomposer_test.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_decomposer_test.cc @@ -32,6 +32,7 @@ limitations under the License. #include "xla/literal.h" #include "xla/literal_util.h" #include "xla/shape.h" +#include "xla/xla_data.pb.h" namespace xla { namespace { From 75515015d234daebc53428a60682db514ea157e0 Mon Sep 17 00:00:00 2001 From: Harsha H S Date: Wed, 23 Sep 2026 07:17:00 -0700 Subject: [PATCH 10/28] PR #48355: [ROCm] Add mi350 backend and replace gfx1250 with mi450 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/48355 πŸ“ Summary of Changes Replace gfx1250 with mi450 and add mi350 test backend πŸš€ Kind of Contribution πŸ“š Documentation, πŸ§ͺ Tests Copybara import of the project: -- 7b2c0050c55bded52b02312ddc8fe0ce8e8e25e0 by Harsha Havanur Shamsundara : Add mi350 backend and replace gfx1250 with mi450 -- 5d0781bfc582a2d31213686c0b7c0c4f7396ccba by Harsha Havanur Shamsundara : Fix errors -- acb1275081c47d5c4883d2d0a8d5cb9ee84bf51f by Harsha Havanur Shamsundara : Address review comments -- 8ed01ad35415530bad0d150644936de5a3e811ab by Harsha HS : Fix lit tests to include mi300 and mi450 gpu configs -- 09e10b8fc26ad401a7ee0965895dc15c2b65e3d7 by Harsha HS : Add mi300 config spec -- 3cb49cc2d1322bf33b7ca1143e8629444c17625d by Harsha HS : Fix HBM3 and device memory size Merging this change closes #48355 PiperOrigin-RevId: 986736971 --- third_party/xla/build_tools/lint/tags.py | 1 + .../gpu/target_config/specs/mi300.txtpb | 40 +++++++++++++++++++ .../specs/{gfx1250.txtpb => mi450.txtpb} | 0 .../gpu/target_config/target_config_map.bzl | 2 +- third_party/xla/xla/backends/gpu/tests/BUILD | 14 +++++-- .../gpu/tests/triton_tdm_elemwise.hlo | 12 ++++-- third_party/xla/xla/lit.bzl | 2 +- .../rocm/rocm_compute_capability.h | 4 +- 8 files changed, 65 insertions(+), 10 deletions(-) create mode 100644 third_party/xla/xla/backends/gpu/target_config/specs/mi300.txtpb rename third_party/xla/xla/backends/gpu/target_config/specs/{gfx1250.txtpb => mi450.txtpb} (100%) diff --git a/third_party/xla/build_tools/lint/tags.py b/third_party/xla/build_tools/lint/tags.py index a909af93153f8e..6da264a0827519 100644 --- a/third_party/xla/build_tools/lint/tags.py +++ b/third_party/xla/build_tools/lint/tags.py @@ -104,6 +104,7 @@ "xla_interpreter": "Uses interpreter backend.", "xla_cpu": "Uses CPU backend.", "xla_amdgpu_any": "Uses ROCm backend.", + "xla_mi350": "Runs on MI350 backend.", "xla_nvgpu_any": "Uses NVIDIA GPU backend.", "xla_intelgpu_any": "Uses Intel GPU backend.", # Below tags are emitted alongside `requires-gpu-x` tags, which is what the diff --git a/third_party/xla/xla/backends/gpu/target_config/specs/mi300.txtpb b/third_party/xla/xla/backends/gpu/target_config/specs/mi300.txtpb new file mode 100644 index 00000000000000..869d97f8f69d01 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/target_config/specs/mi300.txtpb @@ -0,0 +1,40 @@ +# Copyright 2026 The OpenXLA Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +gpu_device_info { + threads_per_block_limit: 1024 + threads_per_warp: 64 + shared_memory_per_block: 65536 + shared_memory_per_block_optin: 65536 + shared_memory_per_core: 65536 + threads_per_core_limit: 2048 + core_count: 304 + fpus_per_core: 128 + block_dim_limit_x: 2147483647 + block_dim_limit_y: 65536 + block_dim_limit_z: 65536 + memory_bandwidth: 5300000000000 + l2_cache_size: 4194304 + clock_rate_ghz: 2.1 + device_memory_size: 206158430208 + registers_per_core_limit: 131072 + registers_per_block_limit: 131072 + rocm_compute_capability { + gcn_arch_name: "gfx942:sramecc+:xnack-" + } +} +platform_name: "ROCM" +dnn_version_info { +} +device_description_str: "MI300X" diff --git a/third_party/xla/xla/backends/gpu/target_config/specs/gfx1250.txtpb b/third_party/xla/xla/backends/gpu/target_config/specs/mi450.txtpb similarity index 100% rename from third_party/xla/xla/backends/gpu/target_config/specs/gfx1250.txtpb rename to third_party/xla/xla/backends/gpu/target_config/specs/mi450.txtpb diff --git a/third_party/xla/xla/backends/gpu/target_config/target_config_map.bzl b/third_party/xla/xla/backends/gpu/target_config/target_config_map.bzl index d658aa86dcef6e..fac887ef139acf 100644 --- a/third_party/xla/xla/backends/gpu/target_config/target_config_map.bzl +++ b/third_party/xla/xla/backends/gpu/target_config/target_config_map.bzl @@ -38,7 +38,7 @@ target_config_map = { "bmg_g21": "//xla/backends/gpu/target_config:specs/bmg_g21.txtpb", "gb200": "//xla/backends/gpu/target_config:specs/gb200.txtpb", "gb300": "//xla/backends/gpu/target_config:specs/gb300.txtpb", - "gfx1250": "//xla/backends/gpu/target_config:specs/gfx1250.txtpb", + "mi450": "//xla/backends/gpu/target_config:specs/mi450.txtpb", "h100_pcie": "//xla/backends/gpu/target_config:specs/h100_pcie.txtpb", "h100_sxm": "//xla/backends/gpu/target_config:specs/h100_sxm.txtpb", "h100": "//xla/backends/gpu/target_config:specs/h100_sxm.txtpb", diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD index 680c99c6556244..7bc17e886098fe 100644 --- a/third_party/xla/xla/backends/gpu/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/tests/BUILD @@ -769,14 +769,21 @@ lit_test_suite_for_gpus( "vectorization.hlo", "pdl.hlo", ], + "mi300": [ + "memory_bandwidth_fusion.hlo", + "dot_bf16.hlo", + "single_instruction.hlo", + "vectorization.hlo", + "pdl.hlo", + ], "mi350": [ "dot_bf16.hlo", "single_instruction.hlo", "vectorization.hlo", "pdl.hlo", ], - "gfx1250": [ - # TODO: re-enable once gfx1250 memory_bandwidth is finalized. + "mi450": [ + # TODO: re-enable once mi450 memory_bandwidth is finalized. "memory_bandwidth_fusion.hlo", "single_instruction.hlo", "vectorization.hlo", @@ -787,10 +794,11 @@ lit_test_suite_for_gpus( "a100_pcie_80", "a6000", "b200", - "gfx1250", "h100_sxm", "mi200", + "mi300", "mi350", + "mi450", ], tags = ["no-oneapi"], tools = [ diff --git a/third_party/xla/xla/backends/gpu/tests/triton_tdm_elemwise.hlo b/third_party/xla/xla/backends/gpu/tests/triton_tdm_elemwise.hlo index 3e21447232700d..22c085a0638500 100644 --- a/third_party/xla/xla/backends/gpu/tests/triton_tdm_elemwise.hlo +++ b/third_party/xla/xla/backends/gpu/tests/triton_tdm_elemwise.hlo @@ -1,18 +1,22 @@ // RUN: hlo-opt %s --platform=gpu --stage=llvm-before-optimizations --xla_gpu_target_config_filename=%S/../target_config/specs/%{GPU}.txtpb | FileCheck %s --check-prefixes=CHECK-%{GPU},CHECK-%{PTX} --allow-unused-prefixes -// Verify that element-wise Triton fusions targeting gfx1250 use TDM +// Verify that element-wise Triton fusions targeting mi450 family use TDM // (Tensor Descriptor Memory) intrinsics instead of pointer arithmetic. // On non-TDM ROCm targets (e.g. mi200), buffer load/store ops are used. // On CUDA targets, the AMD-specific TDM intrinsics must not appear. -// CHECK-gfx1250: @llvm.amdgcn.tensor.load.to.lds -// CHECK-gfx1250: @llvm.amdgcn.s.wait.tensorcnt -// CHECK-gfx1250: @llvm.amdgcn.tensor.store.from.lds +// CHECK-mi450: @llvm.amdgcn.tensor.load.to.lds +// CHECK-mi450: @llvm.amdgcn.s.wait.tensorcnt +// CHECK-mi450: @llvm.amdgcn.tensor.store.from.lds // CHECK-mi200-NOT: @llvm.amdgcn.{{(tensor\.|s\.wait\.tensorcnt)}} // CHECK-mi200: @llvm.amdgcn.raw.ptr.buffer.load // CHECK-mi200-NOT: @llvm.amdgcn.{{(tensor\.|s\.wait\.tensorcnt)}} +// CHECK-mi300-NOT: @llvm.amdgcn.{{(tensor\.|s\.wait\.tensorcnt)}} +// CHECK-mi300: @llvm.amdgcn.raw.ptr.buffer.load +// CHECK-mi300-NOT: @llvm.amdgcn.{{(tensor\.|s\.wait\.tensorcnt)}} + // CHECK-mi350-NOT: @llvm.amdgcn.{{(tensor\.|s\.wait\.tensorcnt)}} // CHECK-mi350: @llvm.amdgcn.raw.ptr.buffer.load // CHECK-mi350-NOT: @llvm.amdgcn.{{(tensor\.|s\.wait\.tensorcnt)}} diff --git a/third_party/xla/xla/lit.bzl b/third_party/xla/xla/lit.bzl index 415c8525921845..d9fa7784bb15c7 100644 --- a/third_party/xla/xla/lit.bzl +++ b/third_party/xla/xla/lit.bzl @@ -223,7 +223,7 @@ def lit_test_suite_for_gpus( # If there are kwargs that need to be passed to only some of the generated # rules, they should be extracted into separate named arguments. - rocm_gpus = ["mi200", "mi350", "gfx1250"] + rocm_gpus = ["mi200", "mi300", "mi350", "mi450"] for gpu in gpus: is_rocm = gpu in rocm_gpus diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_compute_capability.h b/third_party/xla/xla/stream_executor/rocm/rocm_compute_capability.h index 55c85b8d0595cb..1a6cb860c4a128 100644 --- a/third_party/xla/xla/stream_executor/rocm/rocm_compute_capability.h +++ b/third_party/xla/xla/stream_executor/rocm/rocm_compute_capability.h @@ -94,7 +94,9 @@ class RocmComputeCapability { "gfx1030", // RX68xx / RX69xx "gfx1100", // RX7900 "gfx1101", // RX7700 / RX7800 - "gfx1103", "gfx1150", "gfx1151", "gfx1200", "gfx1201", "gfx1250"}; + "gfx1103", "gfx1150", "gfx1151", "gfx1200", "gfx1201", + "gfx1250", // MI450 + }; bool is_supported_gfx_version() const { return IsThisGfxInAnyList(kSupportedGfxVersions); From 7c42bca686bd2bb3ac3e8f74c704cadf309a5de3 Mon Sep 17 00:00:00 2001 From: Enver Kayaaslan Date: Wed, 23 Sep 2026 07:18:59 -0700 Subject: [PATCH 11/28] Thread-safe specialisation cache on hlo evaluator and cache call computation evals from HloConstantFolding. Specialisation cache on hlo evaluator is thread-safe by mutex-locking. PiperOrigin-RevId: 986737880 --- third_party/xla/xla/hlo/evaluator/BUILD | 1 + .../xla/xla/hlo/evaluator/hlo_evaluator.cc | 12 +++--- .../xla/xla/hlo/evaluator/hlo_evaluator.h | 42 ++++++++++++------- .../xla/hlo/evaluator/hlo_evaluator_test.cc | 15 +++---- .../simplifiers/hlo_constant_folding.cc | 3 +- 5 files changed, 40 insertions(+), 33 deletions(-) diff --git a/third_party/xla/xla/hlo/evaluator/BUILD b/third_party/xla/xla/hlo/evaluator/BUILD index bcb0565fd8f01b..2ae54c6ac85335 100644 --- a/third_party/xla/xla/hlo/evaluator/BUILD +++ b/third_party/xla/xla/hlo/evaluator/BUILD @@ -120,6 +120,7 @@ cc_library( "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:cord", "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/synchronization", "@com_google_absl//absl/types:span", "@eigen_archive//:eigen3", "@tsl//tsl/platform:ml_dtypes", diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc index 7200b9f496ed82..69b3cc8cff771c 100644 --- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc +++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc @@ -3672,10 +3672,10 @@ absl::Status HloEvaluator::HandleCall(const HloInstruction* call) { TF_RET_CHECK(specialization_cache_ != nullptr); - const Literal* cached_result = - specialization_cache_->Find(computation, arg_literals); - if (cached_result != nullptr) { - SetEvaluatedLiteralFor(call, cached_result->Clone()); + if (std::optional cached_result = + specialization_cache_->Find(computation, arg_literals); + cached_result.has_value()) { + SetEvaluatedLiteralFor(call, std::move(*cached_result)); return absl::OkStatus(); } @@ -3686,9 +3686,7 @@ absl::Status HloEvaluator::HandleCall(const HloInstruction* call) { ABSL_ASSIGN_OR_RETURN(Literal result, embedded_evaluator->Evaluate(*computation, arg_literals)); - if (specialization_cache_->Find(computation, arg_literals) == nullptr) { - specialization_cache_->Insert(computation, arg_literals, result.Clone()); - } + specialization_cache_->Insert(computation, arg_literals, result.Clone()); SetEvaluatedLiteralFor(call, std::move(result)); return absl::OkStatus(); diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h index 1453825d8a3579..acc6b37b43dfeb 100644 --- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h +++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h @@ -29,6 +29,7 @@ limitations under the License. #include #include +#include "absl/base/thread_annotations.h" #include "absl/container/flat_hash_map.h" #include "absl/container/inlined_vector.h" #include "absl/container/node_hash_map.h" @@ -38,6 +39,7 @@ limitations under the License. #include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" #include "absl/types/span.h" #include "Eigen/Core" #include "xla/array2d.h" @@ -369,13 +371,13 @@ class HloEvaluator : public ConstDfsHloVisitorWithDefault, // See the description on SpecializationKey for more details. struct SpecializationCache { using EntryMap = absl::flat_hash_map; - using iterator = EntryMap::iterator; - using const_iterator = EntryMap::const_iterator; + + mutable absl::Mutex mu; // A deque never relocates existing elements. It means all existing pointers // and references to previously stored Literals remain valid for the // lifetime of the SpecializationCache. - std::deque arg_storage; - EntryMap entries; + std::deque arg_storage ABSL_GUARDED_BY(mu); + EntryMap entries ABSL_GUARDED_BY(mu); SpecializationCache() = default; SpecializationCache(const SpecializationCache&) = delete; @@ -384,31 +386,39 @@ class HloEvaluator : public ConstDfsHloVisitorWithDefault, SpecializationCache& operator=(SpecializationCache&&) = delete; void Clear() { + absl::MutexLock lock(&mu); entries.clear(); arg_storage.clear(); } - size_t size() const { return entries.size(); } - bool empty() const { return entries.empty(); } - iterator begin() { return entries.begin(); } - iterator end() { return entries.end(); } - const_iterator begin() const { return entries.begin(); } - const_iterator end() const { return entries.end(); } - - const Literal* Find(const SpecializationKey& key) const { + size_t size() const { + absl::MutexLock lock(&mu); + return entries.size(); + } + bool empty() const { + absl::MutexLock lock(&mu); + return entries.empty(); + } + + std::optional Find(const SpecializationKey& key) const { + absl::ReaderMutexLock lock(&mu); auto it = entries.find(key); if (it != entries.end()) { - return &it->second; + return it->second.Clone(); } - return nullptr; + return std::nullopt; } - const Literal* Find(const HloComputation* computation, - absl::Span args) const { + std::optional Find(const HloComputation* computation, + absl::Span args) const { return Find(SpecializationKey(computation, args)); } void Insert(const HloComputation* computation, absl::Span args, Literal result) { + absl::MutexLock lock(&mu); + if (entries.contains(SpecializationKey(computation, args))) { + return; + } std::vector slices; slices.reserve(args.size()); for (const Literal* arg : args) { diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc index 7b813bacaeeac8..e8f47905b40997 100644 --- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc +++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc @@ -8724,15 +8724,12 @@ TEST_F(HloEvaluatorTest, NestedCallSharedCache) { const HloComputation* comp_b = m_->GetComputationWithName("comp_b"); ASSERT_NE(comp_b, nullptr); - int comp_b_entries = 0; - for (const auto& [key, val] : *evaluator.specialization_cache()) { - if (key.computation == comp_b) { - ++comp_b_entries; - EXPECT_TRUE( - LiteralTestUtil::Equal(LiteralUtil::CreateR0(15.0f), val)); - } - } - EXPECT_EQ(comp_b_entries, 1); + Literal arg_5 = LiteralUtil::CreateR0(5.0f); + std::optional cached_b = + evaluator.specialization_cache()->Find(comp_b, {&arg_5}); + ASSERT_TRUE(cached_b.has_value()); + EXPECT_TRUE( + LiteralTestUtil::Equal(LiteralUtil::CreateR0(15.0f), *cached_b)); } TEST_F(HloEvaluatorTest, ClearSpecializationCache) { diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc index 003ef2e539c8a3..4f9f4e286bf1af 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc @@ -610,7 +610,8 @@ absl::StatusOr HloConstantFolding::RunImpl( // default case. This retains the behavior from before while loop support in // HloEvaluator and may be revised. auto evaluator = std::make_unique( - /*max_loop_iterations=*/options_.level == Level::kAggressive ? -1 : 0); + /*max_loop_iterations=*/options_.level == Level::kAggressive ? -1 : 0, + /*cache_call_computation_evals=*/true); // fast-path lets us e.g. use Eigen for matmuls. evaluator->set_use_fast_path(true); From 11b054a9dce982b1bd0fdc0561602218796994c2 Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Wed, 23 Sep 2026 07:19:41 -0700 Subject: [PATCH 12/28] [xla] fix nested dynamic_slice algebraic simplification bug Fixes https://github.com/jax-ml/jax/issues/40849 PiperOrigin-RevId: 986738186 --- .../simplifiers/algebraic_simplifier.cc | 11 +++++++-- .../simplifiers/algebraic_simplifier_test.cc | 24 +++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) 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 39f6046fe0549a..97edf369ef7ba6 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc @@ -8070,19 +8070,26 @@ absl::Status AlgebraicSimplifierVisitor::HandleDynamicSlice( } } - // ds(ds(x,id),inner_id) -> ds(x, id + inner_id) + // ds(ds(x, inner_id), id) -> + // ds(x, clamp(0, id, inner_size - outer_size) + + // clamp(0, inner_id, operand_size - inner_size)) if (operand->opcode() == HloOpcode::kDynamicSlice) { ABSL_RETURN_IF_ERROR(dynamic_slice->ReplaceOperandWithDifferentShape( 0, operand->mutable_operand(0))); for (int64_t i = 1; i < dynamic_slice->operand_count(); ++i) { HloInstruction* index = dynamic_slice->mutable_operand(i); + index = index->AddInstruction(HloInstruction::CreateTernary( + index->shape(), HloOpcode::kClamp, MakeScalarLike(index, 0), index, + MakeScalarLike(index, + operand->dynamic_slice_sizes()[i - 1] - + dynamic_slice->dynamic_slice_sizes()[i - 1]))); HloInstruction* inner_index = operand->mutable_operand(i); inner_index = inner_index->AddInstruction(HloInstruction::CreateTernary( inner_index->shape(), HloOpcode::kClamp, MakeScalarLike(inner_index, 0), inner_index, MakeScalarLike(inner_index, operand->operand(0)->shape().dimensions(i - 1) - - dynamic_slice->dynamic_slice_sizes()[i - 1]))); + operand->dynamic_slice_sizes()[i - 1]))); if (inner_index->shape().element_type() != index->shape().element_type()) { inner_index = inner_index->AddInstruction( 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 757244860a7629..b439bf6ed0d202 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 @@ -15063,6 +15063,30 @@ TEST_F(AlgebraicSimplifierTest, DoNotFoldTransposeIntoScatterWhenDisabled) { absl_testing::IsOkAndHolds(false)); } +TEST_F(AlgebraicSimplifierTest, DynamicSliceOfDynamicSlice) { + constexpr absl::string_view hlo_string = R"( + HloModule module + + ENTRY test { + operand = f32[10] parameter(0) + i = s32[] parameter(1) + j = s32[] parameter(2) + inner_ds = f32[4] dynamic-slice(operand, i), dynamic_slice_sizes={4} + ROOT outer_ds = f32[1] dynamic-slice(inner_ds, j), dynamic_slice_sizes={1} + } + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_THAT(simplifier.Run(module.get()), absl_testing::IsOkAndHolds(true)); + EXPECT_THAT(module->entry_computation()->root_instruction(), + GmockMatch(m::DynamicSlice( + m::Parameter(0), + m::Add(m::Clamp(m::ConstantScalar(0), m::Parameter(2), + m::ConstantScalar(3)), + m::Clamp(m::ConstantScalar(0), m::Parameter(1), + m::ConstantScalar(6)))))); +} + } // namespace } // namespace xla // Trivial comment to force rebuild From c8fbd1d1db512944dd8d6bb8105a33f9abb03b49 Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Wed, 23 Sep 2026 07:52:16 -0700 Subject: [PATCH 13/28] Remove Comparison::GetType() and HloCompareInstruction::type() Migrate remaining callers of Comparison::GetType() and HloCompareInstruction::type() (including XlaBuilder::Compare, hlo_to_mhlo, mhlo_to_hlo, and TPU emitters) to use ComparisonOrder and PrimitiveType directly. Also stop setting the legacy comparison_type field in HloCompareInstruction::ToProto. PiperOrigin-RevId: 986754916 --- third_party/xla/xla/comparison_util.cc | 35 ++++--------------- third_party/xla/xla/comparison_util.h | 8 ----- .../xla/xla/hlo/builder/xla_builder.cc | 34 ++++++++++-------- third_party/xla/xla/hlo/builder/xla_builder.h | 10 ++++-- .../xla/xla/hlo/ir/hlo_instructions.cc | 1 - third_party/xla/xla/hlo/ir/hlo_instructions.h | 3 -- .../hlo_to_mhlo/hlo_function_importer.cc | 24 ++++++++----- .../hlo_to_mhlo/hlo_function_importer.h | 4 +-- .../xla/xla/hlo/translate/mhlo_to_hlo/BUILD | 1 + .../translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc | 20 ++++++++--- 10 files changed, 68 insertions(+), 72 deletions(-) diff --git a/third_party/xla/xla/comparison_util.cc b/third_party/xla/xla/comparison_util.cc index 9d7ffcd7012ad7..aed2072aa2d5d5 100644 --- a/third_party/xla/xla/comparison_util.cc +++ b/third_party/xla/xla/comparison_util.cc @@ -223,42 +223,19 @@ Comparison::Order Comparison::DefaultOrdering(PrimitiveType type) { LOG(FATAL) << "Unsupported type: " << PrimitiveType_Name(type); } -namespace { -Comparison::Type ComparisonTypeFromPrimitiveTypeAndOrder( - PrimitiveType type, Comparison::Order order) { - if (primitive_util::IsFloatingPointType(type) || - primitive_util::IsComplexType(type)) { - return order == Comparison::Order::kTotal - ? Comparison::Type::kFloatTotalOrder - : Comparison::Type::kFloat; - } - if (primitive_util::IsSignedIntegralType(type)) { - return Comparison::Type::kSigned; - } - if (primitive_util::IsUnsignedIntegralType(type) || type == PRED) { - return Comparison::Type::kUnsigned; - } - LOG(FATAL) << "Unexpected: " << PrimitiveType_Name(type); -} -} // namespace - Comparison::Comparison(Direction dir, PrimitiveType type, Order order) - : dir_(dir), - primitive_type_(type), - order_(order), - type_(ComparisonTypeFromPrimitiveTypeAndOrder(type, order)) {} + : dir_(dir), primitive_type_(type), order_(order) { + CHECK(primitive_util::IsArrayType(type)) + << "Unsupported type: " << PrimitiveType_Name(type); +} Comparison::Comparison(Direction dir, PrimitiveType type) - : dir_(dir), - primitive_type_(type), - order_(DefaultOrdering(type)), - type_(DefaultComparisonType(type)) {} + : dir_(dir), primitive_type_(type), order_(DefaultOrdering(type)) {} Comparison::Comparison(Direction dir, Type type) : dir_(dir), primitive_type_(DefaultPrimitiveType(type)), - order_(DefaultOrdering(type)), - type_(type) {} + order_(DefaultOrdering(type)) {} Comparison Comparison::Converse() const { return Comparison(xla::Converse(dir_), primitive_type_, order_); diff --git a/third_party/xla/xla/comparison_util.h b/third_party/xla/xla/comparison_util.h index 1b852d21c98040..a626a8d772dcab 100644 --- a/third_party/xla/xla/comparison_util.h +++ b/third_party/xla/xla/comparison_util.h @@ -115,11 +115,6 @@ class Comparison { inline PrimitiveType GetPrimitiveType() const { return primitive_type_; } inline Order GetOrder() const { return order_; } - [[deprecated("Use GetPrimitiveType() and GetOrder()")]] inline Type GetType() - const { - return type_; - } - inline bool IsEq() const { return dir_ == Direction::kEq; } inline bool IsNe() const { return dir_ == Direction::kNe; } inline bool IsGe() const { return dir_ == Direction::kGe; } @@ -233,9 +228,6 @@ class Comparison { const PrimitiveType primitive_type_; // The ordering of the Comparison, e.g., kPartial. const Order order_; - // The Type of the Comparison. This tries to mesh together the ordering and - // the numerical data classification. - [[deprecated]] const Type type_; }; using ComparisonDirection = Comparison::Direction; diff --git a/third_party/xla/xla/hlo/builder/xla_builder.cc b/third_party/xla/xla/hlo/builder/xla_builder.cc index b96c716823cd80..a672fa261726b7 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder.cc +++ b/third_party/xla/xla/hlo/builder/xla_builder.cc @@ -1400,7 +1400,7 @@ absl::StatusOr BroadcastToOutputShapeWithUnbounded( XlaOp XlaBuilder::BinaryOp(HloOpcode binop, XlaOp lhs, XlaOp rhs, absl::Span broadcast_dimensions, std::optional direction, - std::optional type) { + std::optional order) { return ReportErrorOrReturn([&]() -> absl::StatusOr { ABSL_ASSIGN_OR_RETURN(const Shape* lhs_shape, GetShapePtr(lhs)); ABSL_ASSIGN_OR_RETURN(const Shape* rhs_shape, GetShapePtr(rhs)); @@ -1463,10 +1463,10 @@ XlaOp XlaBuilder::BinaryOp(HloOpcode binop, XlaOp lhs, XlaOp rhs, return InvalidArgument( "kCompare expects a ComparisonDirection, but none provided."); } - if (type == std::nullopt) { + if (order == std::nullopt) { return Compare(shape, updated_lhs, updated_rhs, *direction); } - return Compare(shape, updated_lhs, updated_rhs, *direction, *type); + return Compare(shape, updated_lhs, updated_rhs, *direction, *order); } if (direction.has_value()) { @@ -1491,18 +1491,17 @@ absl::StatusOr XlaBuilder::Compare(const Shape& shape, XlaOp lhs, XlaOp rhs, ComparisonDirection direction) { ABSL_ASSIGN_OR_RETURN(auto operand_shape, GetShape(lhs)); - return Compare( - shape, lhs, rhs, direction, - Comparison::DefaultComparisonType(operand_shape.element_type())); + return Compare(shape, lhs, rhs, direction, + Comparison::DefaultOrdering(operand_shape.element_type())); } absl::StatusOr XlaBuilder::Compare(const Shape& shape, XlaOp lhs, XlaOp rhs, ComparisonDirection direction, - Comparison::Type type) { + ComparisonOrder order) { HloInstructionProto instr; instr.set_comparison_direction(ComparisonDirectionToString(direction)); - instr.set_comparison_type(ComparisonTypeToString(type)); + instr.set_comparison_order(ComparisonOrderToShortString(order)); *instr.mutable_shape() = shape.ToProto(); return AddInstruction(std::move(instr), HloOpcode::kCompare, {lhs, rhs}); } @@ -5696,12 +5695,12 @@ static XlaOp CompareTotalOrder(const XlaOp lhs, const XlaOp rhs, return b->ReportErrorOrReturn([&]() -> absl::StatusOr { ABSL_ASSIGN_OR_RETURN(auto operand_shape, b->GetShape(lhs)); auto operand_element_type = operand_shape.element_type(); - auto compare_type = + auto compare_order = primitive_util::IsFloatingPointType(operand_element_type) - ? Comparison::Type::kFloatTotalOrder - : Comparison::DefaultComparisonType(operand_element_type); + ? ComparisonOrder::kTotal + : Comparison::DefaultOrdering(operand_element_type); return Compare(lhs, rhs, broadcast_dimensions, comparison_direction, - compare_type); + compare_order); }); } @@ -5775,9 +5774,16 @@ XlaOp Compare(const XlaOp lhs, const XlaOp rhs, XlaOp Compare(const XlaOp lhs, const XlaOp rhs, absl::Span broadcast_dimensions, - ComparisonDirection direction, Comparison::Type compare_type) { + ComparisonDirection direction, ComparisonOrder order) { return lhs.builder()->BinaryOp(HloOpcode::kCompare, lhs, rhs, - broadcast_dimensions, direction, compare_type); + broadcast_dimensions, direction, order); +} + +XlaOp Compare(const XlaOp lhs, const XlaOp rhs, + absl::Span broadcast_dimensions, + ComparisonDirection direction, Comparison::Type compare_type) { + return Compare(lhs, rhs, broadcast_dimensions, direction, + Comparison::DefaultOrdering(compare_type)); } XlaOp Compare(const XlaOp lhs, const XlaOp rhs, ComparisonDirection direction) { diff --git a/third_party/xla/xla/hlo/builder/xla_builder.h b/third_party/xla/xla/hlo/builder/xla_builder.h index 065be7eb27763b..3c0e937a8c1088 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder.h +++ b/third_party/xla/xla/hlo/builder/xla_builder.h @@ -1254,7 +1254,7 @@ class XlaBuilder { XlaOp BinaryOp(HloOpcode binop, XlaOp lhs, XlaOp rhs, absl::Span broadcast_dimensions, std::optional direction = std::nullopt, - std::optional type = std::nullopt); + std::optional order = std::nullopt); absl::StatusOr Compare(const Shape& shape, XlaOp lhs, XlaOp rhs, ComparisonDirection direction); @@ -1263,7 +1263,7 @@ class XlaBuilder { virtual absl::StatusOr Compare(const Shape& shape, XlaOp lhs, XlaOp rhs, ComparisonDirection direction, - Comparison::Type type); + ComparisonOrder order); // Internal helper method that does the building for an arbitrary binary op // with same ranked operands that doesn't broadcast. @@ -1519,6 +1519,9 @@ class XlaBuilder { friend XlaOp Compare(XlaOp lhs, XlaOp rhs, absl::Span broadcast_dimensions, ComparisonDirection direction); + friend XlaOp Compare(XlaOp lhs, XlaOp rhs, + absl::Span broadcast_dimensions, + ComparisonDirection direction, ComparisonOrder order); friend XlaOp Compare(XlaOp lhs, XlaOp rhs, absl::Span broadcast_dimensions, ComparisonDirection direction, @@ -2574,6 +2577,9 @@ XlaOp LeTotalOrder(XlaOp lhs, XlaOp rhs, // Enqueues a comparison instruction onto the computation (optionally without // broadcast_dimensions for consistency with others). +XlaOp Compare(XlaOp lhs, XlaOp rhs, + absl::Span broadcast_dimensions, + ComparisonDirection direction, ComparisonOrder order); XlaOp Compare(XlaOp lhs, XlaOp rhs, absl::Span broadcast_dimensions, ComparisonDirection direction, Comparison::Type compare_type); diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.cc b/third_party/xla/xla/hlo/ir/hlo_instructions.cc index 9f5c773561d98f..ce9d59cefa5a12 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.cc @@ -961,7 +961,6 @@ void HloCompareInstruction::ToProto(HloInstructionProto* proto) const { ComparisonDirectionToString(compare_.GetDirection())); proto->set_comparison_order( ComparisonOrderToShortString(compare_.GetOrder())); - proto->set_comparison_type(ComparisonTypeToString(compare_.GetType())); } void HloCompareInstruction::PrintExtraAttributesImpl( diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.h b/third_party/xla/xla/hlo/ir/hlo_instructions.h index 10ae4cba26b71d..054e036d666b7d 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.h +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.h @@ -454,9 +454,6 @@ class HloCompareInstruction : public HloInstruction { std::optional order = std::nullopt); ComparisonDirection direction() const { return compare_.GetDirection(); } ComparisonOrder order() const { return compare_.GetOrder(); } - [[deprecated("Use order()")]] Comparison::Type type() const { - return compare_.GetType(); - } const Comparison& comparison() const { return compare_; } void ToProto(HloInstructionProto* proto) const override; diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc index 0f11cbe3d2f722..b180a3292b3748 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc @@ -1294,10 +1294,10 @@ absl::StatusOr HloFunctionImporter::ImportInstructionImpl( case HloOpcode::kCompare: { auto compare = Cast(instruction); attributes.push_back(ConvertComparisonDirection(compare->direction())); - auto default_type = Comparison::DefaultComparisonType( + auto default_order = Comparison::DefaultOrdering( compare->operand(0)->shape().element_type()); - if (compare->type() != default_type) { - attributes.push_back(ConvertComparisonType(compare->type())); + if (compare->order() != default_order) { + attributes.push_back(ConvertComparisonOrder(compare->order())); } return mlir::stablehlo::CompareOp::create(*func_builder, loc, result_type, operands, attributes) @@ -2592,14 +2592,20 @@ mlir::NamedAttribute HloFunctionImporter::ConvertComparisonDirection( .value())); } -mlir::NamedAttribute HloFunctionImporter::ConvertComparisonType( - Comparison::Type type) { +mlir::NamedAttribute HloFunctionImporter::ConvertComparisonOrder( + ComparisonOrder order) { + mlir::stablehlo::ComparisonType type; + switch (order) { + case ComparisonOrder::kPartial: + type = mlir::stablehlo::ComparisonType::FLOAT; + break; + case ComparisonOrder::kTotal: + type = mlir::stablehlo::ComparisonType::TOTALORDER; + break; + } return builder_->getNamedAttr( "compare_type", - mlir::stablehlo::ComparisonTypeAttr::get( - builder_->getContext(), - mlir::stablehlo::symbolizeComparisonType(ComparisonTypeToString(type)) - .value())); + mlir::stablehlo::ComparisonTypeAttr::get(builder_->getContext(), type)); } mlir::DenseIntElementsAttr HloFunctionImporter::ConvertDimensions( diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.h b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.h index 7b34a32e720f59..1116cff866dbb6 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.h +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.h @@ -201,8 +201,8 @@ class HloFunctionImporter { mlir::NamedAttribute ConvertComparisonDirection( ComparisonDirection direction); - // Converts an XLA Comparison::Type to the corresponding MLIR attribute. - mlir::NamedAttribute ConvertComparisonType(Comparison::Type type); + // Converts an XLA ComparisonOrder to the corresponding MLIR attribute. + mlir::NamedAttribute ConvertComparisonOrder(ComparisonOrder order); // Converts an XLA CustomCallSchedule to the corresponding MLIR attribute. mlir::NamedAttribute ConvertCustomCallSchedule(CustomCallSchedule schedule); diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/BUILD b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/BUILD index 627898aa1fc3d3..9d83f381c4bab8 100644 --- a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/BUILD +++ b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/BUILD @@ -261,6 +261,7 @@ cc_library( "//xla/tsl/platform:statusor", "//xla/tsl/platform:types", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc index 1db8afcf410d30..8d225e18d5711f 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 @@ -28,6 +28,7 @@ limitations under the License. #include "mhlo/transforms/passes.h" #include "absl/container/flat_hash_map.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" @@ -2566,10 +2567,21 @@ mlir::LogicalResult ExportXlaOp(mlir::stablehlo::CompareOp op, xla::XlaOp xla_result; if (type_attr && type_attr.getValue() != mlir::stablehlo::ComparisonType::NOTYPE) { - auto type = xla::StringToComparisonType( - stringifyComparisonType(type_attr.getValue()).str()) - .value(); - xla_result = xla::Compare(lhs, rhs, /*broadcast_dimensions=*/{}, dir, type); + xla::ComparisonOrder order; + switch (type_attr.getValue()) { + case mlir::stablehlo::ComparisonType::FLOAT: + order = xla::ComparisonOrder::kPartial; + break; + case mlir::stablehlo::ComparisonType::TOTALORDER: + case mlir::stablehlo::ComparisonType::SIGNED: + case mlir::stablehlo::ComparisonType::UNSIGNED: + order = xla::ComparisonOrder::kTotal; + break; + case mlir::stablehlo::ComparisonType::NOTYPE: + LOG(FATAL) << "Unreachable"; + } + xla_result = + xla::Compare(lhs, rhs, /*broadcast_dimensions=*/{}, dir, order); } else { xla_result = xla::Compare(lhs, rhs, dir); } From 8044c5c35b75e0aa0cbf09219c89eef1e6c78ebd Mon Sep 17 00:00:00 2001 From: Harsha H S Date: Wed, 23 Sep 2026 08:28:14 -0700 Subject: [PATCH 14/28] PR #49177: [ROCm] Remove extrenious arguments like rocm_path passed to host compiler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/49177 πŸ“ Summary of Changes rocm_path and offload_archs compiler arguments were passed to host compiler which fails to compile host side code specifically external dependencies like boringssl or net_zstd while building TF. 🎯 Justification Remove these compiler flags that were passed to host compiler πŸš€ Kind of Contribution: πŸ› Bug Fix Copybara import of the project: -- a65d3df05790142b5cb9cbffa6cbf500d4d9cd1e by Harsha Havanur Shamsundara : Remove extrenious arguments like rocm_path passed to host compiler Merging this change closes #49177 PiperOrigin-RevId: 986776069 --- .../gpus/crosstool/hipcc_cc_toolchain_config.bzl.tpl | 6 ------ 1 file changed, 6 deletions(-) diff --git a/third_party/xla/third_party/gpus/crosstool/hipcc_cc_toolchain_config.bzl.tpl b/third_party/xla/third_party/gpus/crosstool/hipcc_cc_toolchain_config.bzl.tpl index e6c3f6e17bd28d..5451024a08a09d 100644 --- a/third_party/xla/third_party/gpus/crosstool/hipcc_cc_toolchain_config.bzl.tpl +++ b/third_party/xla/third_party/gpus/crosstool/hipcc_cc_toolchain_config.bzl.tpl @@ -1166,12 +1166,6 @@ def _impl(ctx): flag_sets = [ flag_set( actions = all_compile_actions + all_link_actions, - flag_groups = [ - flag_group( - flags = ["--rocm-path=" + _rocm_path] + - ["--offload-arch=" + arch for arch in _hipcc_config.gpu_architectures], - ), - ], ), ], ) From e9aa87061382e67e029c8d37c7aac5193842dc09 Mon Sep 17 00:00:00 2001 From: Viktor Moros Date: Wed, 23 Sep 2026 08:32:02 -0700 Subject: [PATCH 15/28] [XLA Highway] 128-bit CPS load & store primitives PiperOrigin-RevId: 986778048 --- third_party/xla/xla/pjrt/transpose_kernels.h | 25 ++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/third_party/xla/xla/pjrt/transpose_kernels.h b/third_party/xla/xla/pjrt/transpose_kernels.h index da532efae24e83..bee5e081e84903 100644 --- a/third_party/xla/xla/pjrt/transpose_kernels.h +++ b/third_party/xla/xla/pjrt/transpose_kernels.h @@ -200,6 +200,31 @@ HWY_INLINE void CombineRowsCPS(Cont&& cont, Vs... in) { } } +// Stream -> Pack adapter for loads +template +HWY_INLINE void LoadRowsCPS(D d, const char* a, int64_t lda, Cont&& cont, + Vs... vs) { + if constexpr (I == bs) { + std::forward(cont)(vs...); + } else { + auto v = LoadBytes(d, a + lda * I); + LoadRowsCPS(d, a, lda, std::forward(cont), + vs..., v); + } +} +// Pack -> Stream adapter for 128-bit stores +template +HWY_INLINE void StoreVecs128(D d, char* b, int64_t ldb, Vs... v) { + ForEachIndexed( + [&](auto idx_c, auto&& vec) { + constexpr size_t i = decltype(idx_c)::value; + StoreLanes(d, b, ldb, vec, i * stores_per_row, + std::make_index_sequence{}); + }, + v...); +} + } // namespace HWY_NAMESPACE } // namespace xla HWY_AFTER_NAMESPACE(); From f87634eaff1171e67c8d98cb04198cf168fc4c98 Mon Sep 17 00:00:00 2001 From: dlcompilers-infra-bot Date: Wed, 23 Sep 2026 08:54:34 -0700 Subject: [PATCH 16/28] PR #48643: make sure the live range of buffers in async ops span across the entire async duration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/48643 πŸ“ Summary of Changes The live range of buffers in async ops need to reflect the duration of async-start and async-done. There are some corner cases where xla's buffer assignment would re-use buffers wrapped in the async ops for other parallel ops within the async-start/done interval which is incorrect. This could lead to silent data corruption. 🎯 Justification Incorrect live range leads to data corruption. πŸš€ Kind of Contribution πŸ› Bug Fix πŸ§ͺ Unit Tests: added unit test. Copybara import of the project: -- 8b6b61fcb4dbd99a1ad3cbb4a935f58176d96f67 by TJ Xu : make sure the live range of buffers in async ops span across the entire async duration -- d85fd8e108c38dec2b820584f50ab907ff97cb18 by TJ Xu : fix cpu test failure -- 46d8fc531ff617309312f01ff4108209be9335d8 by TJ Xu : address pr comments -- 5a433749da77b3652dab1a307daa9cffcdc60166 by TJ Xu : address comment and fix a crash when computation has multiple users Merging this change closes #48643 PiperOrigin-RevId: 986790742 --- .../xla/xla/hlo/utils/hlo_live_range.cc | 51 ++++- .../xla/xla/hlo/utils/hlo_live_range_test.cc | 212 ++++++++++++++++++ 2 files changed, 261 insertions(+), 2 deletions(-) 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 f3027bb4cff609..b2dfdaca4e8661 100644 --- a/third_party/xla/xla/hlo/utils/hlo_live_range.cc +++ b/third_party/xla/xla/hlo/utils/hlo_live_range.cc @@ -368,11 +368,29 @@ void HloLiveRange::CalculateBufferStartEndMap() { instruction.IsRoot() ? computation_span_times_[computation].end : entry.second; - // If the instruction is in an asynchronous context, extend the live range - // until the end of the async-done instruction. + // If the instruction is in an asynchronous context, adjust its live range + // to cover the async window precisely: + // - end_time is extended to the async-done instruction so the buffer + // remains live for the full async duration. + // - start_time is tightened to the first-fully-bound instruction + // (async-start for standard chains, async-update for late-binding + // chains), because FlattenSchedule inlines the async computation + // immediately before that instruction and the inner buffer is not live + // before it fires. auto async_context_it = computations_in_async_context_.find(computation); if (async_context_it != computations_in_async_context_.end()) { const HloComputation* async_context = async_context_it->second; + // async_context can have multiple callers when the same async-wrapped + // computation is shared by several async-start/update instructions + // (e.g. one outside and one inside a while loop). FlattenSchedule only + // ever inlines the computation once, at whichever caller the schedule + // walk reaches first, so we take the minimum schedule time across all + // first-fully-bound callers below. This is deterministic (unlike + // picking an arbitrary caller from caller_instructions(), which is + // returned in no particular order) and never overshoots the schedule + // position the instructions were actually flattened at, which would + // otherwise push start_time past definition_end_time. + std::optional tightened_start_time; for (const HloInstruction* caller : async_context->caller_instructions()) { if (caller->IsAsynchronous()) { @@ -397,10 +415,39 @@ void HloLiveRange::CalculateBufferStartEndMap() { definition_end_time = std::max( definition_end_time, computation_span_times_[computation].end); } + // Track the earliest first-fully-bound caller's schedule time. + // FlattenSchedule inlines the async computation's instructions + // immediately before the first-fully-bound instruction (async-start + // for standard chains, async-update for late-binding chains) that + // actually triggers the flattening. The inner buffer is not live + // before that point. + absl::StatusOr is_first_fully_bound = + hlo_instruction_utils::async::IsFirstFullyBound(caller); + if (is_first_fully_bound.ok() && *is_first_fully_bound) { + auto first_bound_it = instruction_schedule_.find(caller); + if (first_bound_it != instruction_schedule_.end()) { + tightened_start_time = + tightened_start_time.has_value() + ? std::min(*tightened_start_time, first_bound_it->second) + : first_bound_it->second; + } + } } } + if (tightened_start_time.has_value()) { + // Cap by definition_end_time (computed above): for computations + // reached transitively through a kCall inside the async-wrapped + // computation, FlattenSchedule inlines their bodies even earlier + // than the async-wrapped computation's own instructions, so a + // caller's schedule time can land after such an instruction's own + // (already-correct) end. Capping avoids pushing start past end + // while still tightening whenever it's safe to do so. + start_time = std::min(*tightened_start_time, definition_end_time); + } VLOG(2) << "Setting the definition end time for op in async context: " << definition_end_time; + VLOG(2) << "Setting the definition start time for op in async context: " + << start_time; } for (const HloValue* value : 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 e3bb28c339bff7..c77a682e24a9d7 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 @@ -1202,5 +1202,217 @@ ENTRY entry { EXPECT_FALSE(root_buffers.empty()); EXPECT_TRUE(HloLiveRange::BufferLivesOut(*root_buffers[0], *aa, entry)); } + +TEST_F(HloLiveRangeTest, AsyncInnerBufferStartExtendedToOuterStart) { + const std::string hlo_string = R"hlo( + HloModule AsyncInnerBufferStart, is_scheduled=true + + %async_wrapped (p: f32[4]) -> f32[4] { + %p = f32[4] parameter(0) + ROOT %inner_op = f32[4] negate(%p) + } + + ENTRY %main (a: f32[4]) -> f32[4] { + %a = f32[4] parameter(0) + %outer_op = f32[4] negate(%a) + %async-start = ((f32[4]), f32[4], u32[]) async-start(%outer_op), + calls=%async_wrapped + %async-done = f32[4] async-done(%async-start) + ROOT %result = f32[4] add(%outer_op, %async-done) + } + )hlo"; + + ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(alias_analysis_, + HloAliasAnalysis::Run(module_.get(), &alias_info_)); + ASSERT_OK_AND_ASSIGN(hlo_live_range_, + HloLiveRange::Run(module_->schedule(), *alias_analysis_, + module_->entry_computation())); + CheckSchedule(); + + HloComputation* async_wrapped = + module_->GetComputationWithName("async_wrapped"); + ASSERT_NE(async_wrapped, nullptr); + const HloInstruction* inner_op = + async_wrapped->GetInstructionWithName("inner_op"); + ASSERT_NE(inner_op, nullptr); + + const HloInstruction* async_done = + module_->entry_computation()->GetInstructionWithName("async-done"); + ASSERT_NE(async_done, nullptr); + const HloInstruction* result = + module_->entry_computation()->GetInstructionWithName("result"); + ASSERT_NE(result, nullptr); + + const HloInstruction* async_start = + module_->entry_computation()->GetInstructionWithName("async-start"); + ASSERT_NE(async_start, nullptr); + + auto inner_range = LiveRangeAt(inner_op); + auto async_start_time = + hlo_live_range_->instruction_schedule().at(async_start); + auto async_done_time = hlo_live_range_->instruction_schedule().at(async_done); + auto result_time = hlo_live_range_->instruction_schedule().at(result); + + EXPECT_EQ(inner_range.start, async_start_time); + EXPECT_GT(result_time, async_done_time); + EXPECT_EQ(inner_range.end, result_time); +} + +TEST_F(HloLiveRangeTest, AsyncInnerParameterStartExtendedToOuterStart) { + const std::string hlo_string = R"hlo( + HloModule AsyncInnerParamStart, is_scheduled=true + + %async_wrapped (p: f32[4]) -> f32[4] { + %p = f32[4] parameter(0) + ROOT %inner_op = f32[4] negate(%p) + } + + ENTRY %main (a: f32[4]) -> f32[4] { + %a = f32[4] parameter(0) + %async-start = ((f32[4]), f32[4], u32[]) async-start(%a), + calls=%async_wrapped + %async-done = f32[4] async-done(%async-start) + ROOT %result = f32[4] negate(%async-done) + } + )hlo"; + + ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(alias_analysis_, + HloAliasAnalysis::Run(module_.get(), &alias_info_)); + ASSERT_OK_AND_ASSIGN(hlo_live_range_, + HloLiveRange::Run(module_->schedule(), *alias_analysis_, + module_->entry_computation())); + CheckSchedule(); + + HloComputation* async_wrapped = + module_->GetComputationWithName("async_wrapped"); + ASSERT_NE(async_wrapped, nullptr); + + const HloInstruction* async_start = + module_->entry_computation()->GetInstructionWithName("async-start"); + ASSERT_NE(async_start, nullptr); + const HloInstruction* inner_param = + async_wrapped->GetInstructionWithName("p"); + const HloInstruction* inner_op = + async_wrapped->GetInstructionWithName("inner_op"); + ASSERT_NE(inner_param, nullptr); + ASSERT_NE(inner_op, nullptr); + + auto async_start_time = + hlo_live_range_->instruction_schedule().at(async_start); + + EXPECT_EQ(LiveRangeAt(inner_param).start, 0); + EXPECT_EQ(LiveRangeAt(inner_op).start, async_start_time); +} + +TEST_F(HloLiveRangeTest, AsyncInnerBufferStartIsAsyncStartTime) { + const std::string hlo_string = R"hlo( + HloModule AsyncInnerOverlap, is_scheduled=true + + %async_wrapped (p: f32[4]) -> f32[4] { + %p = f32[4] parameter(0) + ROOT %inner_op = f32[4] negate(%p) + } + + ENTRY %main (a: f32[4]) -> f32[4] { + %a = f32[4] parameter(0) + %temp = f32[4] negate(%a) + %outer = f32[4] negate(%temp) + %async-start = ((f32[4]), f32[4], u32[]) async-start(%outer), + calls=%async_wrapped + %async-done = f32[4] async-done(%async-start) + ROOT %result = f32[4] negate(%async-done) + } + )hlo"; + + ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(alias_analysis_, + HloAliasAnalysis::Run(module_.get(), &alias_info_)); + ASSERT_OK_AND_ASSIGN(hlo_live_range_, + HloLiveRange::Run(module_->schedule(), *alias_analysis_, + module_->entry_computation())); + CheckSchedule(); + + const HloInstruction* async_start = + module_->entry_computation()->GetInstructionWithName("async-start"); + const HloInstruction* temp = + module_->entry_computation()->GetInstructionWithName("temp"); + HloComputation* async_wrapped = + module_->GetComputationWithName("async_wrapped"); + ASSERT_NE(async_start, nullptr); + ASSERT_NE(temp, nullptr); + ASSERT_NE(async_wrapped, nullptr); + + const HloInstruction* inner_op = + async_wrapped->GetInstructionWithName("inner_op"); + ASSERT_NE(inner_op, nullptr); + + auto async_start_time = + hlo_live_range_->instruction_schedule().at(async_start); + auto temp_range = LiveRangeAt(temp); + auto inner_range = LiveRangeAt(inner_op); + + EXPECT_EQ(inner_range.start, async_start_time); + EXPECT_GT(inner_range.start, temp_range.end); +} + +TEST_F(HloLiveRangeTest, AsyncComputationSharedByMultipleCallersUsesMinStart) { + const std::string hlo_string = R"hlo( + HloModule m, is_scheduled=true + + %nested (q: f32[4]) -> f32[4] { + %q = f32[4] parameter(0) + %t = f32[4] negate(%q) + ROOT %n = f32[4] negate(%t) + } + + %async_computation (p: f32[4]) -> f32[4] { + %p = f32[4] parameter(0) + ROOT %c = f32[4] call(%p), to_apply=%nested + } + + %while_cond (wp: ((f32[4]), f32[4], u32[])) -> pred[] { + %wp = ((f32[4]), f32[4], u32[]) parameter(0) + ROOT %cond = pred[] constant(true) + } + + %while_body (bp: ((f32[4]), f32[4], u32[])) -> ((f32[4]), f32[4], u32[]) { + %bp = ((f32[4]), f32[4], u32[]) parameter(0) + %g = f32[4] get-tuple-element(%bp), index=1 + ROOT %s = ((f32[4]), f32[4], u32[]) async-start(%g), calls=%async_computation + } + + ENTRY %main (a: f32[4]) -> ((f32[4]), f32[4], u32[]) { + %a = f32[4] parameter(0) + %init = ((f32[4]), f32[4], u32[]) async-start(%a), calls=%async_computation + ROOT %loop = ((f32[4]), f32[4], u32[]) while(%init), condition=%while_cond, body=%while_body + } + )hlo"; + + ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(alias_analysis_, + HloAliasAnalysis::Run(module_.get(), &alias_info_)); + ASSERT_OK_AND_ASSIGN(hlo_live_range_, + HloLiveRange::Run(module_->schedule(), *alias_analysis_, + module_->entry_computation())); + CheckSchedule(); + + HloComputation* nested = module_->GetComputationWithName("nested"); + ASSERT_NE(nested, nullptr); + const HloInstruction* t = nested->GetInstructionWithName("t"); + ASSERT_NE(t, nullptr); + + const HloInstruction* init = + module_->entry_computation()->GetInstructionWithName("init"); + ASSERT_NE(init, nullptr); + + auto init_time = hlo_live_range_->instruction_schedule().at(init); + auto t_range = LiveRangeAt(t); + + EXPECT_LE(t_range.start, t_range.end); + EXPECT_LT(t_range.start, init_time); +} + } // namespace } // namespace xla From 0697d0a4509c6092c2c94de0ccb460b6c2011ed2 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 23 Sep 2026 08:56:00 -0700 Subject: [PATCH 17/28] Automated Code Change PiperOrigin-RevId: 986791525 --- tensorflow/compiler/mlir/tensorflow/transforms/BUILD | 1 + tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc | 1 + 2 files changed, 2 insertions(+) diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/BUILD b/tensorflow/compiler/mlir/tensorflow/transforms/BUILD index bc47f98a1ba4a3..d19b0e3a8c45a5 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/BUILD +++ b/tensorflow/compiler/mlir/tensorflow/transforms/BUILD @@ -717,6 +717,7 @@ cc_library( "//tensorflow/core/protobuf:for_core_protos_cc", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/log", + "@com_google_absl//absl/log:vlog_is_on", "@com_google_absl//absl/status", "@llvm-project//llvm:Support", "@llvm-project//mlir:FuncDialect", diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc b/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc index d5284e1ae5736b..1fdd6d46ed3296 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/compiler/mlir/tensorflow/transforms/bridge.h" #include "absl/log/log.h" +#include "absl/log/vlog_is_on.h" #include "llvm/ADT/StringRef.h" #include "mlir/IR/BuiltinOps.h" // from @llvm-project #include "mlir/Pass/PassManager.h" // from @llvm-project From f7226c5d8cdf88b9b28fe0a28b448028012802dd Mon Sep 17 00:00:00 2001 From: Viktor Moros Date: Wed, 23 Sep 2026 09:12:37 -0700 Subject: [PATCH 18/28] [XLA Highway] Add 128-bit vector transpose microkernel PiperOrigin-RevId: 986801854 --- third_party/xla/xla/pjrt/transpose_kernels.h | 28 ++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/third_party/xla/xla/pjrt/transpose_kernels.h b/third_party/xla/xla/pjrt/transpose_kernels.h index bee5e081e84903..5d6dcb61f71ee1 100644 --- a/third_party/xla/xla/pjrt/transpose_kernels.h +++ b/third_party/xla/xla/pjrt/transpose_kernels.h @@ -225,6 +225,34 @@ HWY_INLINE void StoreVecs128(D d, char* b, int64_t ldb, Vs... v) { v...); } +// 128-bit (single block) transpose +template +HWY_FLATTEN void TransposeMicroKernel128(const char* HWY_RESTRICT a, + int64_t lda, char* HWY_RESTRICT b, + int64_t ldb) { + using D = hn::Full128; + static constexpr size_t element_size = sizeof(T); + static constexpr size_t row_bytes = element_size * bs; + static_assert(row_bytes <= kBlockBytes); + const D d; + LoadRowsCPS<0, T, bs, D, row_bytes>(d, a, lda, [&](auto... loads) { + CombineRowsCPS( + [&](auto... vecs) { + constexpr size_t kNumVecs = sizeof...(vecs); + constexpr size_t kBytesInMatrix = element_size * bs * bs; + constexpr size_t kStoresPerRow = bs / kNumVecs; + constexpr size_t kElementsPerVec = kBytesInMatrix / (bs * kNumVecs); + InterleaveWithinBlocksCPS( + [&](auto... transposed) { + StoreVecs128(d, b, ldb, + transposed...); + }, + vecs...); + }, + loads...); + }); +} + } // namespace HWY_NAMESPACE } // namespace xla HWY_AFTER_NAMESPACE(); From fed179cf136a7ad5fffcd73f0da1f69e41e92b2f Mon Sep 17 00:00:00 2001 From: Sam Ginzburg Date: Wed, 23 Sep 2026 09:24:56 -0700 Subject: [PATCH 19/28] [XLA:SPMD] Fix cross partition sort and shape invariants in HandleCustomCallTopK PiperOrigin-RevId: 986808802 --- .../xla/service/spmd/custom_call_handler.cc | 68 ++++---- .../xla/service/spmd/spmd_partitioner_test.cc | 152 +++++++++++++++--- 2 files changed, 169 insertions(+), 51 deletions(-) diff --git a/third_party/xla/xla/service/spmd/custom_call_handler.cc b/third_party/xla/xla/service/spmd/custom_call_handler.cc index f6d8472633c488..7e4931ce77b520 100644 --- a/third_party/xla/xla/service/spmd/custom_call_handler.cc +++ b/third_party/xla/xla/service/spmd/custom_call_handler.cc @@ -141,7 +141,7 @@ constexpr char kSPMDOpMultiRotate[] = "_SPMDInternalOp_MultiRotate"; constexpr char kSPMDOpMultiSlice[] = "_SPMDInternalOp_MultiSlice"; constexpr char kSPMDOpWrap[] = "_SPMDInternalOp_Wrap"; -bool IsTopKIndexOutputUsed(HloInstruction* hlo) { +bool IsTopKIndexOutputUsed(const HloInstruction* hlo) { if (hlo->IsRoot()) { return true; } @@ -162,7 +162,14 @@ bool IsTopKIndexOutputUsed(HloInstruction* hlo) { absl::Status SpmdPartitioningVisitor::HandleCustomCallTopK( HloInstruction* hlo) { - if (!hlo->operand(0)->has_sharding()) { + // TopkRewriter flattens rank >= 2 inputs to 2D and leaves rank 1 inputs as + // 1D; this handler assumes a 2D [batch, sort] input. + if (hlo->operand_count() != 1 || !hlo->operand(0)->has_sharding() || + !hlo->operand(0)->shape().IsArray() || + hlo->operand(0)->shape().dimensions().size() != 2 || + !hlo->shape().IsTuple() || hlo->shape().tuple_shapes().size() != 2 || + !hlo->shape().tuple_shapes(0).IsArray() || + hlo->shape().tuple_shapes(0).dimensions().size() != 2) { return DefaultAction(hlo); } @@ -195,19 +202,11 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCallTopK( auto partitioned_input = GetPartitionedHlo(input).PadWithValue( CreateFirstWithType(element_type, &b_)); - auto partition_state = partitioned_input.state(); - auto replicated_sharding = HloSharding::Replicate(); - // If batch dimension is partitioned, partial replicated on sort dimension. - if (batch_dim_partition > 1) { - auto sharding_grouped = - hlo_sharding_util::GroupShardingOnDims(sharding, {batch_dim}); - partition_state = CreatePerGroupPartitioningState( - partitioned_input.state(), sharding_grouped.device_groups, - partitioned_input.state().b); - replicated_sharding = - hlo_sharding_util::PartiallyReplicateTiledShardingOnDims(sharding, - {sort_dim}); - } + const HloSharding replicated_sharding = + batch_dim_partition > 1 + ? hlo_sharding_util::PartiallyReplicateTiledShardingOnDims(sharding, + {sort_dim}) + : HloSharding::Replicate(); // Each partition needs to do TopK separately, thus the base shape // becomes [batch_size, k * shard_count]. @@ -243,7 +242,15 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCallTopK( HloInstruction* slice_sort_value = nullptr; HloInstruction* slice_index_value = nullptr; + const int64_t partitioned_batch_size = + CeilOfRatio(batch_size, batch_dim_partition); if (IsTopKIndexOutputUsed(hlo)) { + const hlo_sharding_util::GroupedSharding sharding_grouped = + hlo_sharding_util::GroupShardingOnDims(sharding, {batch_dim}); + PartitionedHlo::PartitioningState partition_state = + CreatePerGroupPartitioningState(partitioned_input.state(), + sharding_grouped.device_groups, + partitioned_input.state().b); // Get index from TopK. HloInstruction* index_gte = b_.AddInstruction(HloInstruction::CreateGetTupleElement( @@ -280,21 +287,18 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCallTopK( XlaBuilder b("Sort.Compare"); XlaComputation comparator = CreateScalarComparisonComputation( "compare-value-and-index", {input->shape().element_type(), S32}, - {Gt, Lt}, &b); + {GtTotalOrder, LtTotalOrder}, &b); ABSL_ASSIGN_OR_RETURN(HloComputation * compare_computation, XlaComputationToHloComputation(comparator, module_)); // Each partition needs to do TopK separately, thus the base shape for sort // becomes [ceil(batch_size / batch_dim_partition), k * shard_count]. const Shape sort_shape = ShapeUtil::MakeTupleShape( - {ShapeUtil::MakeShape( - hlo->operand(0)->shape().element_type(), - {CeilOfRatio(batch_size, batch_dim_partition), k * shard_count}), - ShapeUtil::MakeShape( - S32, - {CeilOfRatio(batch_size, batch_dim_partition), k * shard_count})}); - auto sort = b_.AddInstruction(HloInstruction::CreateSort( + {ShapeUtil::MakeShape(element_type, + {partitioned_batch_size, k * shard_count}), + ShapeUtil::MakeShape(S32, {partitioned_batch_size, k * shard_count})}); + HloInstruction* sort = b_.AddInstruction(HloInstruction::CreateSort( sort_shape, sort_dim, {replicated_value_gte, replicated_index_gte}, - compare_computation, true)); + compare_computation, /*is_stable=*/false)); sort->set_sharding( replicated_sharding.GetTupleSharding(sort->shape()).value()); PartitionedHlo replicated_sort(sort, replicated_shape, @@ -317,20 +321,19 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCallTopK( // Sort only the replicated values. XlaBuilder b("Sort.Compare"); XlaComputation comparator = CreateScalarComparisonComputation( - "compare-value", {input->shape().element_type()}, {Gt}, &b); + "compare-value", {input->shape().element_type()}, {GtTotalOrder}, &b); ABSL_ASSIGN_OR_RETURN(HloComputation * compare_computation, XlaComputationToHloComputation(comparator, module_)); const Shape sort_shape = ShapeUtil::MakeShape( - hlo->operand(0)->shape().element_type(), - {CeilOfRatio(batch_size, batch_dim_partition), k * shard_count}); - auto sort = b_.AddInstruction( + element_type, {partitioned_batch_size, k * shard_count}); + HloInstruction* sort = b_.AddInstruction( HloInstruction::CreateSort(sort_shape, sort_dim, {replicated_value_gte}, - compare_computation, true)); + compare_computation, /*is_stable=*/false)); sort->set_sharding(replicated_sharding); - const Shape sort_replicated_shape = ShapeUtil::MakeShape( - hlo->operand(0)->shape().element_type(), {batch_size, k * shard_count}); + const Shape sort_replicated_shape = + ShapeUtil::MakeShape(element_type, {batch_size, k * shard_count}); PartitionedHlo replicated_sort(sort, sort_replicated_shape, MakePartitioningState()); @@ -342,7 +345,8 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCallTopK( HloInstruction* dummy_index_constant = b_.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))); slice_index_value = b_.AddInstruction(HloInstruction::CreateBroadcast( - ShapeUtil::MakeShape(S32, {batch_size, k}), dummy_index_constant, {})); + ShapeUtil::MakeShape(S32, {partitioned_batch_size, k}), + dummy_index_constant, {})); slice_index_value->set_sharding(replicated_sharding); } diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc index e7d679a29789b2..024a14599209d8 100644 --- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc +++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc @@ -3749,9 +3749,13 @@ ENTRY %cluster_2013453984438090939__.47 ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/2)); VLOG(1) << module->ToString(); - auto custom_call = FindInstruction(module.get(), "custom-call.1"); + const HloInstruction* custom_call = + FindInstruction(module.get(), "custom-call.1"); + ASSERT_NE(custom_call, nullptr); EXPECT_EQ(custom_call->operand(0)->shape().dimensions(1), 104832); - auto sort = FindInstruction(module.get(), "sort"); + const HloInstruction* sort = FindInstruction(module.get(), "sort"); + ASSERT_NE(sort, nullptr); + EXPECT_FALSE(Cast(sort)->is_stable()); EXPECT_EQ(sort->operand(0)->shape().dimensions(1), 4000); EXPECT_EQ(sort->operand(1)->shape().dimensions(1), 4000); } @@ -3779,9 +3783,13 @@ ENTRY entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/8)); - auto custom_call = FindInstruction(module.get(), "custom-call.1"); + const HloInstruction* custom_call = + FindInstruction(module.get(), "custom-call.1"); + ASSERT_NE(custom_call, nullptr); EXPECT_EQ(custom_call->operand(0)->shape().dimensions(1), 32128); - auto sort = FindInstruction(module.get(), "sort"); + const HloInstruction* sort = FindInstruction(module.get(), "sort"); + ASSERT_NE(sort, nullptr); + EXPECT_FALSE(Cast(sort)->is_stable()); EXPECT_EQ(sort->operand(0)->shape().dimensions(0), 1); EXPECT_EQ(sort->operand(0)->shape().dimensions(1), 2); EXPECT_EQ(sort->operand(1)->shape().dimensions(0), 1); @@ -3811,9 +3819,13 @@ ENTRY entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/8)); VLOG(1) << module->ToString(); - auto custom_call = FindInstruction(module.get(), "custom-call.1"); + const HloInstruction* custom_call = + FindInstruction(module.get(), "custom-call.1"); + ASSERT_NE(custom_call, nullptr); EXPECT_EQ(custom_call->operand(0)->shape().dimensions(1), 16064); - auto sort = FindInstruction(module.get(), "sort"); + const HloInstruction* sort = FindInstruction(module.get(), "sort"); + ASSERT_NE(sort, nullptr); + EXPECT_FALSE(Cast(sort)->is_stable()); EXPECT_EQ(sort->operand(0)->shape().dimensions(0), 2); EXPECT_EQ(sort->operand(0)->shape().dimensions(1), 4); EXPECT_EQ(sort->operand(1)->shape().dimensions(0), 2); @@ -17470,10 +17482,19 @@ ENTRY %entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/2)); VLOG(1) << module->ToString(); - auto sort_instruction = FindInstruction(module.get(), HloOpcode::kSort); + const HloInstruction* sort_instruction = + FindInstruction(module.get(), HloOpcode::kSort); + ASSERT_NE(sort_instruction, nullptr); EXPECT_THAT(sort_instruction, op::Shape("bf16[64,80]{1,0}")); - auto topk_instruction = FindInstruction(module.get(), HloOpcode::kCustomCall); - auto topk_operand = topk_instruction->operand(0); + EXPECT_FALSE(Cast(sort_instruction)->is_stable()); + EXPECT_EQ(Cast( + sort_instruction->to_apply()->root_instruction()) + ->order(), + Comparison::Order::kTotal); + const HloInstruction* topk_instruction = + FindInstruction(module.get(), HloOpcode::kCustomCall); + ASSERT_NE(topk_instruction, nullptr); + const HloInstruction* topk_operand = topk_instruction->operand(0); EXPECT_EQ(topk_instruction->custom_call_target(), "TopK"); EXPECT_THAT(topk_instruction, op::Shape("(bf16[64,40]{1,0}, s32[64,40]{1,0})")); @@ -17503,10 +17524,21 @@ ENTRY %entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/2)); VLOG(1) << module->ToString(); - auto sort_instruction = FindInstruction(module.get(), HloOpcode::kSort); - CHECK_NE(sort_instruction, nullptr); - auto topk_instruction = FindInstruction(module.get(), HloOpcode::kCustomCall); - auto topk_operand = topk_instruction->operand(0); + const HloInstruction* sort_instruction = + FindInstruction(module.get(), HloOpcode::kSort); + ASSERT_NE(sort_instruction, nullptr); + EXPECT_FALSE(Cast(sort_instruction)->is_stable()); + for (const HloInstruction* inst : + sort_instruction->to_apply()->instructions()) { + if (inst->opcode() == HloOpcode::kCompare) { + EXPECT_EQ(Cast(inst)->order(), + Comparison::Order::kTotal); + } + } + const HloInstruction* topk_instruction = + FindInstruction(module.get(), HloOpcode::kCustomCall); + ASSERT_NE(topk_instruction, nullptr); + const HloInstruction* topk_operand = topk_instruction->operand(0); EXPECT_EQ(topk_instruction->custom_call_target(), "TopK"); EXPECT_THAT(topk_instruction, op::Shape("(bf16[32,40]{1,0}, s32[32,40]{1,0})")); @@ -17584,10 +17616,19 @@ ENTRY %entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/2)); VLOG(1) << module->ToString(); - auto sort_instruction = FindInstruction(module.get(), HloOpcode::kSort); + const HloInstruction* sort_instruction = + FindInstruction(module.get(), HloOpcode::kSort); + ASSERT_NE(sort_instruction, nullptr); EXPECT_THAT(sort_instruction, op::Shape("bf16[32,40]{1,0}")); - auto topk_instruction = FindInstruction(module.get(), HloOpcode::kCustomCall); - auto topk_operand = topk_instruction->operand(0); + EXPECT_FALSE(Cast(sort_instruction)->is_stable()); + EXPECT_EQ(Cast( + sort_instruction->to_apply()->root_instruction()) + ->order(), + Comparison::Order::kTotal); + const HloInstruction* topk_instruction = + FindInstruction(module.get(), HloOpcode::kCustomCall); + ASSERT_NE(topk_instruction, nullptr); + const HloInstruction* topk_operand = topk_instruction->operand(0); EXPECT_EQ(topk_instruction->custom_call_target(), "TopK"); EXPECT_THAT(topk_instruction, op::Shape("(bf16[32,40]{1,0}, s32[32,40]{1,0})")); @@ -17616,16 +17657,89 @@ ENTRY %entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/2)); VLOG(1) << module->ToString(); - auto sort_instruction = FindInstruction(module.get(), HloOpcode::kSort); + const HloInstruction* sort_instruction = + FindInstruction(module.get(), HloOpcode::kSort); + ASSERT_NE(sort_instruction, nullptr); EXPECT_THAT(sort_instruction, op::Shape("bf16[64,80]{1,0}")); - auto topk_instruction = FindInstruction(module.get(), HloOpcode::kCustomCall); - auto topk_operand = topk_instruction->operand(0); + EXPECT_FALSE(Cast(sort_instruction)->is_stable()); + EXPECT_EQ(Cast( + sort_instruction->to_apply()->root_instruction()) + ->order(), + Comparison::Order::kTotal); + const HloInstruction* topk_instruction = + FindInstruction(module.get(), HloOpcode::kCustomCall); + ASSERT_NE(topk_instruction, nullptr); + const HloInstruction* topk_operand = topk_instruction->operand(0); EXPECT_EQ(topk_instruction->custom_call_target(), "TopK"); EXPECT_THAT(topk_instruction, op::Shape("(bf16[64,40]{1,0}, s32[64,40]{1,0})")); EXPECT_THAT(topk_operand, op::Shape("bf16[64,128000]{1,0}")); } +TEST_P(SpmdPartitioningTest, TopKCustomCallTwoDimsShardedValuesOnly) { + constexpr absl::string_view kHloString = R"( +HloModule module + +region_695.22546 { + Arg_2.22549 = s32[] parameter(2) + Arg_3.22550 = s32[] parameter(3) + Arg_0.22547 = bf16[] parameter(0) + Arg_1.22548 = bf16[] parameter(1) + ROOT compare.22551 = pred[] compare(Arg_0.22547, Arg_1.22548), direction=GT, order=TOTAL +} + +ENTRY %entry { + %multiply.43401 = bf16[64,256000]{1,0} parameter(0), sharding={devices=[2,2]<=[4]} + %custom-call = (bf16[64,40]{1,0}, s32[64,40]{1,0}) custom-call(bf16[64,256000]{1,0} %multiply.43401), custom_call_target="TopK", called_computations={%region_695.22546}, sharding={{devices=[2,2]<=[4]}, {devices=[2,2]<=[4]}} + ROOT %get-tuple-element.336 = bf16[64,40]{1,0} get-tuple-element((bf16[64,40]{1,0}, s32[64,40]{1,0}) %custom-call), index=0 +})"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + PartitionComputation(kHloString, /*num_devices=*/4)); + const HloInstruction* sort_instruction = + FindInstruction(module.get(), HloOpcode::kSort); + ASSERT_NE(sort_instruction, nullptr); + EXPECT_THAT(sort_instruction, op::Shape("bf16[32,80]{1,0}")); + EXPECT_FALSE(Cast(sort_instruction)->is_stable()); + EXPECT_EQ(Cast( + sort_instruction->to_apply()->root_instruction()) + ->order(), + Comparison::Order::kTotal); + const HloInstruction* topk_instruction = + FindInstruction(module.get(), HloOpcode::kCustomCall); + ASSERT_NE(topk_instruction, nullptr); + EXPECT_THAT(topk_instruction, + op::Shape("(bf16[32,40]{1,0}, s32[32,40]{1,0})")); + EXPECT_THAT(topk_instruction->operand(0), op::Shape("bf16[32,128000]{1,0}")); +} + +TEST_P(SpmdPartitioningTest, TopKCustomCallRank1) { + constexpr absl::string_view kHloString = R"( +HloModule module + +region_695.22546 { + Arg_2.22549 = s32[] parameter(2) + Arg_3.22550 = s32[] parameter(3) + Arg_0.22547 = bf16[] parameter(0) + Arg_1.22548 = bf16[] parameter(1) + ROOT compare.22551 = pred[] compare(Arg_0.22547, Arg_1.22548), direction=GT, order=TOTAL +} + +ENTRY %entry { + %p0 = bf16[256000]{0} parameter(0), sharding={devices=[2]<=[2]} + %custom-call = (bf16[40]{0}, s32[40]{0}) custom-call(bf16[256000]{0} %p0), custom_call_target="TopK", called_computations={%region_695.22546}, sharding={{replicated}, {replicated}} + ROOT %get-tuple-element.336 = bf16[40]{0} get-tuple-element((bf16[40]{0}, s32[40]{0}) %custom-call), index=0 +})"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + PartitionComputation(kHloString, /*num_devices=*/2)); + const HloInstruction* topk_instruction = + FindInstruction(module.get(), HloOpcode::kCustomCall); + ASSERT_NE(topk_instruction, nullptr); + EXPECT_THAT(topk_instruction, op::Shape("(bf16[40]{0}, s32[40]{0})")); + EXPECT_THAT(topk_instruction->operand(0), op::Shape("bf16[256000]{0}")); +} + TEST_P(SpmdPartitioningTest, TopKCustomCallManualSharding) { absl::string_view hlo_string = R"( HloModule module From 067900221579c11c85a82c229eed4698e513cdcd Mon Sep 17 00:00:00 2001 From: Viktor Moros Date: Wed, 23 Sep 2026 10:26:55 -0700 Subject: [PATCH 20/28] [XLA Highway] Add 256-bit square vector transpose microkernel PiperOrigin-RevId: 986848921 --- third_party/xla/xla/pjrt/transpose_kernels.h | 82 ++++++++++++++++++++ 1 file changed, 82 insertions(+) diff --git a/third_party/xla/xla/pjrt/transpose_kernels.h b/third_party/xla/xla/pjrt/transpose_kernels.h index 5d6dcb61f71ee1..59f6b37361b5b2 100644 --- a/third_party/xla/xla/pjrt/transpose_kernels.h +++ b/third_party/xla/xla/pjrt/transpose_kernels.h @@ -253,6 +253,88 @@ HWY_FLATTEN void TransposeMicroKernel128(const char* HWY_RESTRICT a, }); } +// 256-bit (two block) transposes +#if HWY_MIN_BYTES >= 32 +template +HWY_INLINE void StoreRows256(D d, char* b, int64_t ldb, Vs... v) { + ForEachIndexed( + [&](auto idx_c, auto&& vec) { + constexpr size_t i = decltype(idx_c)::value; + hn::StoreU(vec, d, reinterpret_cast*>(b + ldb * i)); + }, + v...); +} + +// Invokes a continuation with an explicit parameter pack, rearranging the +// arguments such that all even-indexed elements precede all odd-indexed +// elements. +template +HWY_INLINE void CallContWithSeparatedImpl(Cont&& cont, + std::index_sequence, + Vs&&... vs) { + std::forward(cont)(GetFromPack<2 * I>(std::forward(vs)...)..., + GetFromPack<2 * I + 1>(std::forward(vs)...)...); +} + +template +HWY_INLINE void CallContWithSeparated(Cont&& cont, Vs&&... vs) { + CallContWithSeparatedImpl(std::forward(cont), + std::make_index_sequence{}, + std::forward(vs)...); +} + +// Loads a square block of memory into 256-bit vectors, partitioning the +// data spatially. Invokes the continuation with the vectors corresponding +// to the left half of the block followed by the right half. +template +HWY_INLINE void BuildVecs256SquareCPS(D d, DHalf d_half, const char* a, + int64_t lda, Cont& cont, Vs... vs) { + if constexpr (I == bs / 2) { + CallContWithSeparated(cont, vs...); + } else { + auto* row0 = reinterpret_cast*>(a + lda * I); + auto row1 = + reinterpret_cast*>(a + lda * (I + bs / 2)); + auto row0_lo = hn::LoadU(d_half, row0); + auto row0_hi = hn::LoadU(d_half, row0 + kBlockBytes / sizeof(*row0)); + auto row1_lo = hn::LoadU(d_half, row1); + auto row1_hi = hn::LoadU(d_half, row1 + kBlockBytes / sizeof(*row0)); + auto v_lo = hn::Combine(d, row1_lo, row0_lo); + auto v_hi = hn::Combine(d, row1_hi, row0_hi); + BuildVecs256SquareCPS(d, d_half, a, lda, cont, vs..., v_lo, + v_hi); + } +} + +template +HWY_INLINE void BuildVecs256SquareCPS(D d, DHalf d_half, const char* a, + int64_t lda, Cont& cont) { + BuildVecs256SquareCPS<0, T, bs>(d, d_half, a, lda, cont); +} + +template +HWY_FLATTEN void TransposeMicroKernel256Square(const char* HWY_RESTRICT a, + int64_t lda, + char* HWY_RESTRICT b, + int64_t ldb) { + using D = hn::FixedTag; + using DHalf = hn::Half; + static constexpr size_t element_size = sizeof(T); + static constexpr size_t row_bytes = element_size * bs; + static_assert(row_bytes == 2 * kBlockBytes); + static_assert(bs % 2 == 0); + const D d; + const DHalf d_half; + auto cont = [&](auto... vecs) { + InterleaveWithinBlocksCPS( + [&](auto... transposed) { StoreRows256(d, b, ldb, transposed...); }, + vecs...); + }; + BuildVecs256SquareCPS(d, d_half, a, lda, cont); +} +#endif // HWY_MIN_BYTES >= 32 + } // namespace HWY_NAMESPACE } // namespace xla HWY_AFTER_NAMESPACE(); From 4ee259160078a8890b2e65729cf070859980eecb Mon Sep 17 00:00:00 2001 From: Christian Sigg Date: Wed, 23 Sep 2026 10:40:48 -0700 Subject: [PATCH 21/28] Update CUDA Tile to v13.4.0. PiperOrigin-RevId: 986858417 --- .../xla/third_party/cuda_tile/cuda_tile.BUILD | 35 ++++ .../cuda_tile/patches/constructor.patch | 2 +- .../cuda_tile/patches/pow_op_alias.patch | 13 ++ .../xla/third_party/cuda_tile/patches/series | 3 +- .../patches/symbol_op_interface.patch | 163 ------------------ .../cuda_tile/patches/unused_variable.patch | 58 +++++++ .../xla/third_party/cuda_tile/workspace.bzl | 4 +- 7 files changed, 111 insertions(+), 167 deletions(-) create mode 100644 third_party/xla/third_party/cuda_tile/patches/pow_op_alias.patch delete mode 100644 third_party/xla/third_party/cuda_tile/patches/symbol_op_interface.patch create mode 100644 third_party/xla/third_party/cuda_tile/patches/unused_variable.patch diff --git a/third_party/xla/third_party/cuda_tile/cuda_tile.BUILD b/third_party/xla/third_party/cuda_tile/cuda_tile.BUILD index 72a4d67e2d0297..a2765e58eeb66a 100644 --- a/third_party/xla/third_party/cuda_tile/cuda_tile.BUILD +++ b/third_party/xla/third_party/cuda_tile/cuda_tile.BUILD @@ -156,6 +156,35 @@ gentbl_cc_library( deps = [":CudaTileTdFiles"], ) +gentbl_cc_library( + name = "CudaTileAttrCodeGenIncGen", + tbl_outs = [ + ( + ["-gen-enum-attr-list"], + "include/cuda_tile/Dialect/CudaTile/IR/CudaTileEnumAttrs.inc", + ), + ( + ["-gen-gpu-arch-def"], + "include/cuda_tile/Dialect/CudaTile/IR/GpuArchitectures.inc", + ), + ( + ["-gen-gpu-arch-constants"], + "include/cuda_tile/Dialect/CudaTile/IR/GpuArchConstants.inc", + ), + ( + ["-gen-hint-key-accessors"], + "include/cuda_tile/Dialect/CudaTile/IR/HintKeyAccessors.inc", + ), + ( + ["-gen-hint-key-impl"], + "include/cuda_tile/Dialect/CudaTile/IR/HintKeyImpl.inc", + ), + ], + tblgen = ":cuda-tile-tblgen", + td_file = "include/cuda_tile/Dialect/CudaTile/IR/AttrDefs.td", + deps = [":CudaTileTdFiles"], +) + gentbl_cc_library( name = "CudaTileRemarksIncGen", tbl_outs = [ @@ -219,12 +248,15 @@ cc_library( ], visibility = ["//visibility:public"], deps = [ + ":CudaTileAttrCodeGenIncGen", ":CudaTileAttrDefsIncGen", ":CudaTileDialectIncGen", ":CudaTileInterfacesIncGen", ":CudaTileOpsCanonicalizationIncGen", ":CudaTileOpsIncGen", + ":CudaTileRemarksIncGen", ":CudaTileTypesIncGen", + "@llvm-project//llvm:Core", "@llvm-project//llvm:Support", "@llvm-project//mlir:ArithDialect", "@llvm-project//mlir:BytecodeOpInterface", @@ -276,6 +308,7 @@ cc_library( "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", + "@llvm-project//mlir:SideEffectInterfaces", "@llvm-project//mlir:TransformUtils", ], ) @@ -363,7 +396,9 @@ cc_library( ":CudaTileBytecodeTypeIncGen", ":CudaTileDialect", "@llvm-project//llvm:Support", + "@llvm-project//mlir:FunctionInterfaces", "@llvm-project//mlir:IR", + "@llvm-project//mlir:Support", "@llvm-project//mlir:TranslateLib", ], ) diff --git a/third_party/xla/third_party/cuda_tile/patches/constructor.patch b/third_party/xla/third_party/cuda_tile/patches/constructor.patch index 98b074b4be7358..06374ab761ad3d 100644 --- a/third_party/xla/third_party/cuda_tile/patches/constructor.patch +++ b/third_party/xla/third_party/cuda_tile/patches/constructor.patch @@ -3,7 +3,7 @@ Remove redundant defaulted copy/move constructors to preserve aggregate status. diff --git a/tools/cuda-tile-tblgen/CudaTileOp.h b/tools/cuda-tile-tblgen/CudaTileOp.h --- a/tools/cuda-tile-tblgen/CudaTileOp.h +++ b/tools/cuda-tile-tblgen/CudaTileOp.h -@@ -55,11 +55,6 @@ +@@ -56,11 +56,6 @@ std::optional> selectedVariants; ParameterType type; diff --git a/third_party/xla/third_party/cuda_tile/patches/pow_op_alias.patch b/third_party/xla/third_party/cuda_tile/patches/pow_op_alias.patch new file mode 100644 index 00000000000000..14bb426cb6d73b --- /dev/null +++ b/third_party/xla/third_party/cuda_tile/patches/pow_op_alias.patch @@ -0,0 +1,13 @@ +diff --git a/include/cuda_tile/Dialect/CudaTile/IR/Ops.h b/include/cuda_tile/Dialect/CudaTile/IR/Ops.h +--- a/include/cuda_tile/Dialect/CudaTile/IR/Ops.h ++++ b/include/cuda_tile/Dialect/CudaTile/IR/Ops.h +@@ -59,6 +59,9 @@ + + namespace mlir::cuda_tile { + ++/// Backwards compatibility alias for PowOp which was renamed to FPowFOp in v13.4.0. ++using PowOp = FPowFOp; ++ + // Helper function to extract cuda_tile::ModuleOp + cuda_tile::ModuleOp extractCudaTileModuleOp(Operation *op); + diff --git a/third_party/xla/third_party/cuda_tile/patches/series b/third_party/xla/third_party/cuda_tile/patches/series index f26d3afb3502f8..9134fcfde4c7bd 100644 --- a/third_party/xla/third_party/cuda_tile/patches/series +++ b/third_party/xla/third_party/cuda_tile/patches/series @@ -1,2 +1,3 @@ constructor.patch -symbol_op_interface.patch +unused_variable.patch +pow_op_alias.patch diff --git a/third_party/xla/third_party/cuda_tile/patches/symbol_op_interface.patch b/third_party/xla/third_party/cuda_tile/patches/symbol_op_interface.patch deleted file mode 100644 index d48d6ab5b4408a..00000000000000 --- a/third_party/xla/third_party/cuda_tile/patches/symbol_op_interface.patch +++ /dev/null @@ -1,163 +0,0 @@ -Fix SymbolOpInterface usages for LLVM r632326. - -diff --git a/include/cuda_tile/Dialect/CudaTile/IR/Ops.td b/include/cuda_tile/Dialect/CudaTile/IR/Ops.td ---- a/include/cuda_tile/Dialect/CudaTile/IR/Ops.td -+++ b/include/cuda_tile/Dialect/CudaTile/IR/Ops.td -@@ -2178,7 +2178,7 @@ def CudaTile_FToIOp : CudaTileConversion - //===----------------------------------------------------------------------===// - - def CudaTile_EntryOp : CudaTileCoreOpDef<"entry", "13.1", [ -- FunctionOpInterface, IsolatedFromAbove, OpAsmOpInterface, SingleBlock, -+ SymbolName, FunctionOpInterface, IsolatedFromAbove, OpAsmOpInterface, SingleBlock, - SingleBlockImplicitTerminator<"ReturnOp"> - ]> { - let summary = "Define a tile kernel"; -@@ -2217,6 +2217,12 @@ def CudaTile_EntryOp : CudaTileCoreOpDef - /// Returns the result types of this function. - ArrayRef getResultTypes() { return getFunctionType().getResults(); } - -+ // SymbolOpInterface Methods -+ mlir::SymbolTable::Visibility getVisibility() { -+ return mlir::SymbolTable::Visibility::Public; -+ } -+ void setVisibility(mlir::SymbolTable::Visibility vis) {} -+ - static void build(::mlir::OpBuilder &odsBuilder, - ::mlir::OperationState &odsState, - ::mlir::StringAttr sym_name, -@@ -2350,7 +2356,7 @@ def CudaTile_GetGlobalOp : CudaTileCore - // GlobalOp - //===----------------------------------------------------------------------===// - --def CudaTile_GlobalOp : CudaTileCoreOpDef<"global", "13.1", [Symbol]> { -+def CudaTile_GlobalOp : CudaTileCoreOpDef<"global", "13.1", [Symbol, SymbolName]> { - let summary = "Allocate static global memory"; - - let description = [{ -@@ -2404,6 +2410,23 @@ def CudaTile_GlobalOp : CudaTileCoreOpDe - - let assemblyFormat = "custom($symbol_visibility) (`constant` $constant^)? $sym_name (`alignment` `=` $alignment^)? attr-dict custom($value)"; - let hasVerifier = 1; -+ -+ let extraClassDeclaration = [{ -+ // SymbolOpInterface Methods -+ mlir::SymbolTable::Visibility getVisibility() { -+ if (getSymbolVisibility() == SymbolVisibility::Public) { -+ return mlir::SymbolTable::Visibility::Public; -+ } -+ return mlir::SymbolTable::Visibility::Private; -+ } -+ void setVisibility(mlir::SymbolTable::Visibility vis) { -+ if (vis == mlir::SymbolTable::Visibility::Public) { -+ setSymbolVisibility(SymbolVisibility::Public); -+ } else { -+ setSymbolVisibility(SymbolVisibility::Private); -+ } -+ } -+ }]; - } - - //===----------------------------------------------------------------------===// -diff --git a/include/cuda_tile/Dialect/CudaTile/IR/SharedFuncParserAndPrinter.h b/include/cuda_tile/Dialect/CudaTile/IR/SharedFuncParserAndPrinter.h ---- a/include/cuda_tile/Dialect/CudaTile/IR/SharedFuncParserAndPrinter.h -+++ b/include/cuda_tile/Dialect/CudaTile/IR/SharedFuncParserAndPrinter.h -@@ -24,9 +24,9 @@ template - ParseResult parseFuncOp(OpAsmParser &parser, OperationState &result) { - // Parse the name as a symbol. - StringAttr nameAttr; -- if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(), -- result.attributes)) -+ if (parser.parseSymbolName(nameAttr, "sym_name", result.attributes)) { - return failure(); -+ } - - // Parse the function signature using custom parsing that supports both - // short form (tile>) and long form (!cuda_tile.tile>) types -@@ -80,10 +80,9 @@ ParseResult parseFuncOp(OpAsmParser &par - template - void printFuncOp(OpTy op, OpAsmPrinter &printer) { - // Print the operation and the function name. -- auto funcName = -- op.getOperation() -- ->template getAttrOfType(SymbolTable::getSymbolAttrName()) -- .getValue(); -+ auto funcName = op.getOperation() -+ ->template getAttrOfType("sym_name") -+ .getValue(); - printer << ' '; - printer.printSymbolName(funcName); - auto fnType = op.getFunctionType(); -@@ -95,8 +94,8 @@ void printFuncOp(OpTy op, OpAsmPrinter & - /*printEmptyBlock=*/false); - printer.printOptionalAttrDict( - op.getOperation()->getAttrs(), -- {op.getArgAttrsAttrName(), op.getFunctionTypeAttrName(), -- SymbolTable::getSymbolAttrName(), op.getResAttrsAttrName()}); -+ {op.getArgAttrsAttrName(), op.getFunctionTypeAttrName(), "sym_name", -+ op.getResAttrsAttrName()}); - } - - } // end namespace mlir::cuda_tile. -diff --git a/include/cuda_tile/Dialect/CudaTile/IR/TestingOps.td b/include/cuda_tile/Dialect/CudaTile/IR/TestingOps.td ---- a/include/cuda_tile/Dialect/CudaTile/IR/TestingOps.td -+++ b/include/cuda_tile/Dialect/CudaTile/IR/TestingOps.td -@@ -31,7 +31,7 @@ include "cuda_tile/Dialect/CudaTile/IR/T - //===----------------------------------------------------------------------===// - - def CudaTile_Test_FuncOp : CudaTileTestingOpDef<"func", "250.0", [ -- IsolatedFromAbove, FunctionOpInterface, SingleBlock, OpAsmOpInterface, -+ IsolatedFromAbove, SymbolName, FunctionOpInterface, SingleBlock, OpAsmOpInterface, - SingleBlockImplicitTerminator<"ReturnOp"> - ]> { - -@@ -56,6 +56,12 @@ def CudaTile_Test_FuncOp : CudaTileTesti - - /// Returns the result types of this function. - ArrayRef getResultTypes() { return getFunctionType().getResults(); } -+ -+ // SymbolOpInterface Methods -+ mlir::SymbolTable::Visibility getVisibility() { -+ return mlir::SymbolTable::Visibility::Public; -+ } -+ void setVisibility(mlir::SymbolTable::Visibility vis) {} - }]; - } - -diff --git a/lib/Dialect/CudaTile/IR/CudaTile.cpp b/lib/Dialect/CudaTile/IR/CudaTile.cpp ---- a/lib/Dialect/CudaTile/IR/CudaTile.cpp -+++ b/lib/Dialect/CudaTile/IR/CudaTile.cpp -@@ -2669,9 +2669,9 @@ constexpr char kOptimizationHintsAttr[] - ParseResult EntryOp::parse(OpAsmParser &parser, OperationState &result) { - // Parse the name as a symbol. - StringAttr nameAttr; -- if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(), -- result.attributes)) -+ if (parser.parseSymbolName(nameAttr, "sym_name", result.attributes)) { - return failure(); -+ } - - // Parse the function signature using custom parsing that supports both - // short form (tile>) and long form (!cuda_tile.tile>) types -@@ -2736,9 +2736,7 @@ ParseResult EntryOp::parse(OpAsmParser & - void EntryOp::print(OpAsmPrinter &printer) { - // Print the operation and the function name. - auto funcName = -- getOperation() -- ->getAttrOfType(SymbolTable::getSymbolAttrName()) -- .getValue(); -+ getOperation()->getAttrOfType("sym_name").getValue(); - printer << ' '; - printer.printSymbolName(funcName); - auto fnType = getFunctionType(); -@@ -2754,9 +2752,8 @@ void EntryOp::print(OpAsmPrinter &printe - /*printEmptyBlock=*/false); - printer.printOptionalAttrDict( - getOperation()->getAttrs(), -- {getArgAttrsAttrName(), getFunctionTypeAttrName(), -- SymbolTable::getSymbolAttrName(), getResAttrsAttrName(), -- kOptimizationHintsAttr}); -+ {getArgAttrsAttrName(), getFunctionTypeAttrName(), "sym_name", -+ getResAttrsAttrName(), kOptimizationHintsAttr}); - } - - LogicalResult EntryOp::verify() { diff --git a/third_party/xla/third_party/cuda_tile/patches/unused_variable.patch b/third_party/xla/third_party/cuda_tile/patches/unused_variable.patch new file mode 100644 index 00000000000000..3258c90744e083 --- /dev/null +++ b/third_party/xla/third_party/cuda_tile/patches/unused_variable.patch @@ -0,0 +1,58 @@ +diff --git a/tools/cuda-tile-tblgen/BytecodeTypeAnalysis.cpp b/tools/cuda-tile-tblgen/BytecodeTypeAnalysis.cpp +--- a/tools/cuda-tile-tblgen/BytecodeTypeAnalysis.cpp ++++ b/tools/cuda-tile-tblgen/BytecodeTypeAnalysis.cpp +@@ -68,7 +68,7 @@ + return Kind::OptionalEnum; + } + +- if (const auto *defInit = dyn_cast_if_present(param.getDef())) { ++ if (const auto* defInit = dyn_cast_if_present(param.getDef())) { + if (defInit->getDef()->isSubClassOf("CudaTileAttrInterfaceParam")) { + return Kind::PolymorphicNestedAttr; + } +@@ -110,7 +110,6 @@ + cppType(param.getCppType().str()), + cppStorageType(param.getCppStorageType().str()), + isOptional(param.isOptional()), kind(classifyParameter(param)) { +- + // Extract version from wrapped parameters (CudaTileTypeParam adds + // sinceVersion). Use getValue() to check if field exists before accessing - + // raw parameters like ArrayRefParameter don't have this field in their class +@@ -153,8 +152,6 @@ + "::" + typeDef.getCppClassName().str()), + typeTagValue(tagValue), sinceVersion(version.str()), + needsReverseOrder(typeDef.getCppClassName() == "TileType") { +- const Record *typeRecord = typeDef.getDef(); +- + // Analyze and validate all parameters. + for (const auto &attrParam : typeDef.getParameters()) { + BytecodeTypeParameter param(attrParam); +@@ -239,7 +236,7 @@ + // through it. Both qualified forms ("ns::Foo" and "::ns::Foo") are indexed + // so lookups don't depend on how OptionalParameter spells the name. + llvm::StringSet<> enumAttrCppTypes; +- for (const Record *r : records.getAllDerivedDefinitions("CudaTileEnumAttr")) { ++ for (const Record* r : records.getAllDerivedDefinitions("CudaTileEnumAttr")) { + AttrOrTypeDef def(r); + std::string qualified = + (def.getDialect().getCppNamespace() + "::" + def.getCppClassName()) +@@ -250,7 +247,7 @@ + + // Build map of CudaTileTypeDef for matching. + StringMap cudaTileTypeDefRecords; +- for (const Record *typeRecord : ++ for (const Record* typeRecord : + records.getAllDerivedDefinitions("CudaTileTypeDef")) { + cudaTileTypeDefRecords[AttrOrTypeDef(typeRecord).getCppClassName()] = + typeRecord; +@@ -327,8 +324,8 @@ + + // Validate every OptionalEnum-classified parameter is actually a + // CudaTileEnumAttr-derived enum-attribute wrapper. +- for (const CudaTileType &type : structure.cudaTileTypes) { +- for (const BytecodeTypeParameter &p : type.parameters) { ++ for (const CudaTileType& type : structure.cudaTileTypes) { ++ for (const BytecodeTypeParameter& p : type.parameters) { + if (p.kind != BytecodeTypeParameter::Kind::OptionalEnum) { + continue; + } diff --git a/third_party/xla/third_party/cuda_tile/workspace.bzl b/third_party/xla/third_party/cuda_tile/workspace.bzl index 72f49935a0ca06..c04963da4a4130 100644 --- a/third_party/xla/third_party/cuda_tile/workspace.bzl +++ b/third_party/xla/third_party/cuda_tile/workspace.bzl @@ -5,8 +5,8 @@ load("//third_party:repo.bzl", "tf_mirror_urls") def repo(): """Imports CUDA Tile IR.""" - CUDA_TILE_COMMIT = "af2417041cc939b87ef56d92cfdcf61737c5457e" - CUDA_TILE_SHA256 = "81597e49469171bf8fa7319fbd44ebe133001521f484589e3dd3fb3fad282dc0" + CUDA_TILE_COMMIT = "7e8e2e68fa219716103824c01f7303367cf7df8d" + CUDA_TILE_SHA256 = "2398f88dd0dbcda3732934c8527485c3941d3a6453c98b6c249461e6309ea44c" quilt_http_archive( name = "cuda_tile", From e8d9ba5babd223697689af5495641cd2cac374af Mon Sep 17 00:00:00 2001 From: Naums Mogers Date: Wed, 23 Sep 2026 10:49:23 -0700 Subject: [PATCH 22/28] [Mosaic] Add single vreg vector load with expanded semantics Add `tpu.expand_load_vreg`, the physical (single vreg) expanded load, as the mirror of `tpu.compress_store_vreg`. The mask is mandatory and one dimensional over the lanes. Packed types keep a 2-D `(lane_count, packing)` vreg, so only dimension 0 has to match the mask. PiperOrigin-RevId: 986863897 --- .../xla/xla/mosaic/dialect/tpu/tpu_ops.cc | 25 ++++++ .../xla/xla/mosaic/dialect/tpu/tpu_ops.td | 36 ++++++++ .../dialect/tpu/tpu_ops_verification_test.cc | 86 ++++++++++++++++++- 3 files changed, 146 insertions(+), 1 deletion(-) diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc index 70fbf7739effb6..89a0171789b02c 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc @@ -1214,6 +1214,31 @@ LogicalResult verifyLoadOp(Op op) { return success(); } +LogicalResult ExpandLoadVregOp::verify() { + MemRefType ref_ty = getBase().getType(); + if (ref_ty.getMemorySpace() && !HasMemorySpace(ref_ty, MemorySpace::kVmem)) { + return emitOpError("Expected base memref to be in VMEM."); + } + VectorType result_ty = getResult().getType(); + if (result_ty.getElementType() != ref_ty.getElementType()) { + return emitOpError("Expected base and result element type to match"); + } + if (llvm::size(getIndices()) != ref_ty.getRank()) { + return emitOpError("Expected ") << ref_ty.getRank() << " indices."; + } + // Note: We deliberately do not go through verifyLoadOp, which rejects masked + // loads of non-32-bit element types. The mask here is mandatory and narrow + // element types are allowed. They are rejected later, in the vector layout + // passes. + VectorType mask_ty = getMask().getType(); + if (result_ty.getShape()[0] != mask_ty.getShape()[0]) { + return emitOpError( + "Expected result dimension 0 to match mask dimension 0: ") + << result_ty.getShape()[0] << " vs " << mask_ty.getShape()[0] << "."; + } + return success(); +} + LogicalResult VectorLoadOp::verify() { const MemRefType ref_ty = getBase().getType(); if (llvm::size(getIndices()) != ref_ty.getRank()) { diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td index 3d050637762cd5..084deae550db80 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td @@ -377,6 +377,42 @@ def TPU_VectorCompressStoreOp : TPU_Op<"vector_compress_store", [DefaultMemWrite let hasVerifier = 1; } +def TPU_ExpandLoadVregOp : TPU_Op<"expand_load_vreg", [DefaultMemRead]> { + let summary = "Performs a masked, expanded load of a single vreg."; + let description = [{ + `tpu.expand_load_vreg` performs a masked, expanded load of a single vreg + from memory starting at the address given by `base[indices]`. It is the + physical (vreg-level) counterpart of `tpu.vector_expand_load`. + + Masked-out lanes do not consume an address: the enabled lanes are filled + from contiguous, sequential locations in memory starting at + `base[indices]`. The n-th enabled lane reads `base[indices] + n`. + + For example, loading with mask [true, false, true, false] yields + [`base[indices]`, , `base[indices] + 1`, ]. + + Elements in masked-out lanes are undefined. + + base : Memref to load from. + indices: Scalar indices into `base`, identifying the position of the first + loaded element. The length must be the same as the rank of `base`. + mask : 1-D lane mask. + result : The loaded vreg. For packed (sub-32-bit) types, the vreg is 2-D + with shape `(lane_count, packing)`, and a lane loads all of its + `packing` subelements together. + }]; + let arguments = (ins + AnyMemRef:$base, + Variadic:$indices, + VectorOfRankAndType<[1], [I1]>:$mask + ); + let results = (outs TPU_Vreg:$result); + let assemblyFormat = [{ + $base `[` $indices `]` `,` $mask attr-dict `:` type($base) `,` type($result) `,` type($mask) + }]; + let hasVerifier = 1; +} + // tpu.vector_load loads a vector from memory into a register. // // base : Memref to load from. diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc index c6b31587e3bd02..2fa8f8df4ecae9 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc @@ -377,7 +377,7 @@ TEST_F(TpuOpsVerificationTest, CompressStoreVregValid) { TEST_F(TpuOpsVerificationTest, CompressStoreVregValid2D) { auto c0 = Create(0); Value memref = AllocaI32({8, 2}, MemorySpace::kVmem); - Value vector_to_store = ConstantI32Vector(/*shape=*/{8, 2}, /*values=*/{1}); + Value vector_to_store = ConstantI32Vector(/*shape=*/{8}, /*values=*/{1}); Value mask = ConstantI1Vector( /*shape=*/{8}, /*values=*/{true, false, true, false, true, false, true, false}); @@ -452,6 +452,90 @@ TEST_F(TpuOpsVerificationTest, CompressStoreVregMismatchedMaskShape) { "Expected valueToStore dimension 0 to match mask dimension 0"))); } +TEST_F(TpuOpsVerificationTest, ExpandLoadVregValid) { + auto c0 = Create(0); + Value memref = AllocaI32({8}, MemorySpace::kVmem); + Value mask = ConstantI1Vector( + /*shape=*/{8}, + /*values=*/{true, false, true, false, true, false, true, false}); + auto el = Create( + /*result=*/VectorType::get({8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0}, + /*mask=*/mask); + + ASSERT_OK(VerifyOp(el)); +} + +// Packed types use a (lane_count, packing) vreg with a lane mask, so only +// dimension 0 has to match the mask. +TEST_F(TpuOpsVerificationTest, ExpandLoadVregValidPacked) { + auto c0 = Create(0); + Value memref = + Create( + GetMemRefType({16}, builder().getBF16Type(), MemorySpace::kVmem)) + .getMemref(); + Value mask = ConstantI1Vector( + /*shape=*/{8}, + /*values=*/{true, false, true, false, true, false, true, false}); + auto el = Create( + /*result=*/VectorType::get({8, 2}, builder().getBF16Type()), + /*base=*/memref, + /*indices=*/ValueRange{c0}, + /*mask=*/mask); + + ASSERT_OK(VerifyOp(el)); +} + +TEST_F(TpuOpsVerificationTest, ExpandLoadVregInvalidMemorySpace) { + auto c0 = Create(0); + Value memref = AllocaI32({8}, MemorySpace::kHbm); + Value mask = ConstantI1Vector( + /*shape=*/{8}, + /*values=*/{true, false, true, false, true, false, true, false}); + auto el = Create( + /*result=*/VectorType::get({8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0}, + /*mask=*/mask); + + ASSERT_THAT(VerifyOp(el), + StatusIs(_, HasSubstr("Expected base memref to be in VMEM."))); +} + +TEST_F(TpuOpsVerificationTest, ExpandLoadVregMismatchedMaskShape) { + auto c0 = Create(0); + Value memref = AllocaI32({8}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{4}, + /*values=*/{true, true, true, true}); + auto el = Create( + /*result=*/VectorType::get({8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0}, + /*mask=*/mask); + + ASSERT_THAT( + VerifyOp(el), + StatusIs(_, + HasSubstr("Expected result dimension 0 to match mask dimension " + "0"))); +} + +TEST_F(TpuOpsVerificationTest, ExpandLoadVregMismatchedIndicesCount) { + auto c0 = Create(0); + Value memref = AllocaI32({8}, MemorySpace::kVmem); + Value mask = ConstantI1Vector( + /*shape=*/{8}, + /*values=*/{true, false, true, false, true, false, true, false}); + auto el = Create( + /*result=*/VectorType::get({8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask); + + ASSERT_THAT(VerifyOp(el), StatusIs(_, HasSubstr("Expected 1 indices."))); +} + TEST_F(TpuOpsVerificationTest, VectorCompressStoreValid) { auto c0 = Create(0); Value memref = AllocaI32({4, 16}, MemorySpace::kVmem); From 183a07fedd57318bfd4757af27a6df640e2f212d Mon Sep 17 00:00:00 2001 From: David Pizzuto Date: Wed, 23 Sep 2026 10:55:16 -0700 Subject: [PATCH 23/28] mlir_to_hlo: Check the return value from the Riegeli writer. MLIR's serialization APIs can't check for I/O errors themselves because raw_ostream doesn't have an API for it. So we have to do it ourselves. PiperOrigin-RevId: 986867518 --- third_party/xla/xla/pjrt/BUILD | 1 + third_party/xla/xla/pjrt/mlir_to_hlo.cc | 17 +++++++++++ third_party/xla/xla/pjrt/mlir_to_hlo_test.cc | 31 ++++++++++++++++++++ 3 files changed, 49 insertions(+) diff --git a/third_party/xla/xla/pjrt/BUILD b/third_party/xla/xla/pjrt/BUILD index 832ba429b85b6c..bc2859c3567a6b 100644 --- a/third_party/xla/xla/pjrt/BUILD +++ b/third_party/xla/xla/pjrt/BUILD @@ -848,6 +848,7 @@ xla_cc_test( "@com_google_googletest//:gtest_main", "@llvm-project//mlir:IR", "@riegeli//riegeli/bytes:string_writer", + "@riegeli//riegeli/bytes:writer", "@stablehlo//:stablehlo_portable_api", ], ) diff --git a/third_party/xla/xla/pjrt/mlir_to_hlo.cc b/third_party/xla/xla/pjrt/mlir_to_hlo.cc index 2fc7f207516248..6e43b8742321df 100644 --- a/third_party/xla/xla/pjrt/mlir_to_hlo.cc +++ b/third_party/xla/xla/pjrt/mlir_to_hlo.cc @@ -321,6 +321,11 @@ absl::Status SerializeToRiegeli(mlir::ModuleOp mlir_module, return absl::InvalidArgumentError("Writer cannot be null."); } + // If we are already in a failed state, no need to do anything else. + if (!writer->ok()) { + return writer->status(); + } + mlir::MLIRContext* context = mlir_module->getContext(); mlir::BaseScopedDiagnosticHandler diagnostic_handler(context); @@ -408,11 +413,23 @@ absl::Status SerializeToRiegeli(mlir::ModuleOp mlir_module, if (mlir::failed(mlir::stablehlo::serializePortableArtifact( mlir_module, target, os, /*allowOtherDialects=*/allow_mixed_serialization))) { + // If we have an I/O error, return it specifically before we return the + // generic serialization error. + if (!writer->ok()) { + return writer->status(); + } const absl::Status status = diagnostic_handler.ConsumeStatus(); return absl::InvalidArgumentError( absl::StrCat("Failed to serialize StableHLO to plugin version ", target, ";\n\nDetailed error from MLIR: ", status.message())); } + + // The raw_ostream APIs don't have support for reporting I/O errors, so we + // have to check them ourselves. + if (!writer->ok()) { + return writer->status(); + } + return absl::OkStatus(); } diff --git a/third_party/xla/xla/pjrt/mlir_to_hlo_test.cc b/third_party/xla/xla/pjrt/mlir_to_hlo_test.cc index 63d66f04e5611b..6a1a365c962846 100644 --- a/third_party/xla/xla/pjrt/mlir_to_hlo_test.cc +++ b/third_party/xla/xla/pjrt/mlir_to_hlo_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "xla/pjrt/mlir_to_hlo.h" +#include #include #include @@ -27,6 +28,7 @@ limitations under the License. #include "mlir/IR/MLIRContext.h" #include "mlir/IR/OwningOpRef.h" #include "riegeli/bytes/string_writer.h" +#include "riegeli/bytes/writer.h" #include "stablehlo/api/PortableApi.h" #include "xla/hlo/testlib/test.h" #include "xla/tsl/platform/statusor.h" @@ -284,5 +286,34 @@ TEST(MlirToHloTest, SerializeToRiegeliTest) { EXPECT_THAT(buffer, IsVhloArtifact("1.0.0")); } +class FailingWriter : public riegeli::Writer { + protected: + bool PushSlow(size_t, size_t) override { + Fail(absl::FailedPreconditionError("We were always going to fail.")); + return false; + } +}; + +TEST(MlirToHloTest, SerializeToRiegeliReportsIOErrors) { + constexpr absl::string_view kProgram = + R"( + func.func @add(%arg0: tensor<1x2xf32>) -> tensor<1x2xf32> { + %cst = stablehlo.constant dense<1.0> : tensor<1x2xf32> + %0 = stablehlo.add %arg0, %cst : tensor<1x2xf32> + return %0 : tensor<1x2xf32> + } + )"; + mlir::MLIRContext context; + ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef module, + ParseMlirModuleString(kProgram, context)); + + FailingWriter writer; + EXPECT_THAT(SerializeToRiegeli(*module, /*requested_target=*/"1.0.0", + /*sdy_version=*/"0.0.1", /*inplace=*/false, + /*allow_mixed_serialization=*/true, &writer), + StatusIs(absl::StatusCode::kFailedPrecondition, + HasSubstr("We were always going to fail."))); +} + } // namespace } // namespace xla From 703b689940ac640dd2862e95ed0fc76ae86668d2 Mon Sep 17 00:00:00 2001 From: Mikhail Goncharov Date: Wed, 23 Sep 2026 11:07:36 -0700 Subject: [PATCH 24/28] [XLA:GPU] disable RTX6000PRO test for cross compile on gb200 they have different compute capabilities. Previously that was not an issue but new tiling is able to emit module with triton and autotuner tries to compile the model and fails with "Compute capability mismatch; expected {12.0}, but was {10.0a}" PiperOrigin-RevId: 986876297 --- .../xla/xla/service/gpu/gb200_cross_compilation_test.cc | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/third_party/xla/xla/service/gpu/gb200_cross_compilation_test.cc b/third_party/xla/xla/service/gpu/gb200_cross_compilation_test.cc index 61d8fed3dc4275..5b32a62424680e 100644 --- a/third_party/xla/xla/service/gpu/gb200_cross_compilation_test.cc +++ b/third_party/xla/xla/service/gpu/gb200_cross_compilation_test.cc @@ -109,7 +109,9 @@ TEST_F(Gb200CrossCompilationTest, CrossCompilationToB200) { EXPECT_TRUE(!AutotunerCache::ResultCacheIsEmpty() || !entries.empty()); } -TEST_F(Gb200CrossCompilationTest, CrossCompilationToRTX6000PRO) { +// TODO(b/565341904): RTX6000PRO has compute capability 12, while gb200 +// only 10.a. Test used to pass as we didn't do any autotuning. +TEST_F(Gb200CrossCompilationTest, DISABLED_CrossCompilationToRTX6000PRO) { ASSERT_OK_AND_ASSIGN(std::unique_ptr result, CrossCompileTo(GpuModel::RTX6000PRO)); // Verify that the auto tuner ran. From ce313315f6b854b569d3a40dfd5a5daa1de4ba72 Mon Sep 17 00:00:00 2001 From: Zviki Nozadze Date: Wed, 23 Sep 2026 11:14:32 -0700 Subject: [PATCH 25/28] [XLA] HloDCE eliminates dead async collectives with reverse postorder. PiperOrigin-RevId: 986881086 --- .../xla/hlo/transforms/simplifiers/hlo_dce.cc | 53 +++-- .../transforms/simplifiers/hlo_dce_test.cc | 211 ++++++++++++++++++ 2 files changed, 249 insertions(+), 15 deletions(-) diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_dce.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_dce.cc index c68f5cdf75aff8..5732a50d4a5456 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_dce.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_dce.cc @@ -228,7 +228,8 @@ bool CanRemoveInstruction( bool remove_cross_partition_collective_ops, const std::function(const HloComputation*)>& computation_callers) { - if (!instruction->IsDead() || HasDisableWhileLoopDceAttr(instruction)) { + if (instruction->parent() == nullptr || !instruction->IsDead() || + HasDisableWhileLoopDceAttr(instruction)) { return false; } @@ -250,18 +251,42 @@ bool CanRemoveInstruction( instruction->operand(0)->user_count() != 1)) { return false; } + auto has_dce_side_effect_attr = [](const HloInstruction* inst, + absl::string_view value) { + if (inst == nullptr) { + return false; + } + auto it = + inst->frontend_attributes().map().find(kDceSideEffectFrontendAttribute); + return it != inst->frontend_attributes().map().end() && it->second == value; + }; + // Reverse postorder visits async-done before async-start. If "false" is + // set on async-start, async-done must inspect async_chain_start(). + // otherwise only async-done will be removed, leaving a dangling async-start + if (has_dce_side_effect_attr(instruction, "false") || + has_dce_side_effect_attr(instruction->async_chain_start(), "false")) { + return false; + } if (instruction->HasSideEffect()) { - auto maybe_collective_op = DynCast(instruction); + auto maybe_collective_op = DynCast( + instruction->async_wrapped_instruction() + ? instruction->async_wrapped_instruction() + : instruction); bool allow_collective = remove_cross_partition_collective_ops && maybe_collective_op && !maybe_collective_op->constrain_layout(); bool allow_while = IsRemovableWhile(instruction, remove_cross_partition_collective_ops); - bool allow_custom_call = instruction->IsCustomCall("tpu_custom_call") && - instruction->frontend_attributes().map().contains( - kDceSideEffectFrontendAttribute) && - instruction->frontend_attributes().map().at( - kDceSideEffectFrontendAttribute) == "true"; + // Reverse postorder visits async-done before async-start. If "true" is + // set on async-start, async-done must inspect + // async_chain_start(). otherwise async-done is kept + // alive, which keeps async-start used (user_count > 0) and prevents + // either from being removed. + bool allow_custom_call = + (instruction->IsCustomCall("tpu_custom_call") || + instruction->IsAsynchronous()) && + (has_dce_side_effect_attr(instruction, "true") || + has_dce_side_effect_attr(instruction->async_chain_start(), "true")); if (!allow_collective && !allow_while && !allow_custom_call) { return false; } @@ -274,27 +299,25 @@ absl::StatusOr RemoveDeadRoots( const std::function(const HloComputation*)>& computation_callers) { bool changed = false; - std::vector dead_roots; - for (auto* instruction : computation->instructions()) { + auto post_order = computation->MakeInstructionPostOrder(); + for (auto it = post_order.rbegin(); it != post_order.rend(); ++it) { + HloInstruction* instruction = *it; if (!CanRemoveInstruction(instruction, remove_cross_partition_collective_ops, computation_callers)) { continue; } - dead_roots.push_back(instruction); - } - - for (HloInstruction* dead_root : dead_roots) { - VLOG(1) << "Removing dead root " << dead_root->ToString() + VLOG(1) << "Removing dead root " << instruction->ToString() << " and its unused operands"; ABSL_RETURN_IF_ERROR(computation->RemoveInstructionAndUnusedOperands( - dead_root, /*cleanup=*/std::nullopt, + instruction, /*cleanup=*/std::nullopt, /*ignore_control_dependencies=*/false, /*computation_callers=*/computation_callers)); changed = true; } return changed; } + absl::Status RemoveDeadParametersFromEntryComputationLayout( HloModule* module, std::vector& dead_parameter_indexes) { if (dead_parameter_indexes.empty()) { diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_dce_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_dce_test.cc index 3646f62f1864dd..fe8911e33aed6c 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_dce_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_dce_test.cc @@ -1227,5 +1227,216 @@ ENTRY main { EXPECT_FALSE(changed); } +TEST_F(HloDceTest, DceAlreadyRemovedInstructionInReversePostOrder) { + constexpr absl::string_view kHloString = R"hlo( +HloModule test_already_removed_guard + +ENTRY main { + p0 = f32[8] parameter(0) + dead_pure_producer = f32[8] negate(p0) + dead_consumer = f32[8] custom-call(dead_pure_producer), custom_call_target="tpu_custom_call", custom_call_has_side_effect=true, frontend_attributes={xla_allow_dce_side_effecting_op="true"} + ROOT root = f32[8] copy(p0) +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloString)); + HloDCE dce; + // Removing `dead_consumer` cascades removal to its pure operand + // `dead_pure_producer` (detaching it from parent). Subsequent reverse + // postorder iteration must skip `dead_pure_producer` rather than crashing. + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_TRUE(changed); + EXPECT_EQ(FindInstruction(module.get(), "dead_consumer"), nullptr); + EXPECT_EQ(FindInstruction(module.get(), "dead_pure_producer"), nullptr); +} + +TEST_F(HloDceTest, PreserveInstructionWithFrontendAttributeFalse) { + constexpr absl::string_view kHloString = R"hlo( +HloModule test_preserve_self_attr_false + +ENTRY main { + p0 = f32[8] parameter(0) + live_ag = f32[16] all-gather(p0), dimensions={0}, replica_groups={{0,1}}, channel_id=1, frontend_attributes={xla_allow_dce_side_effecting_op="false"} + ROOT root = f32[8] copy(p0) +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloString)); + HloDCE dce(/*remove_cross_partition_collective_ops=*/true); + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_FALSE(changed); + EXPECT_NE(FindInstruction(module.get(), "live_ag"), nullptr); +} + +TEST_F(HloDceTest, PreserveAsyncPairWhenStartHasFrontendAttributeFalse) { + constexpr absl::string_view kHloString = R"hlo( +HloModule test_preserve_async_pair_start_attr_false + +async_comp { + p0 = f32[8] parameter(0) + ROOT ag = f32[16] all-gather(p0), dimensions={0}, replica_groups={{0,1}}, channel_id=1 +} + +ENTRY main { + p0 = f32[8] parameter(0) + live_start = ((f32[8]), f32[16]) async-start(p0), calls=async_comp, frontend_attributes={xla_allow_dce_side_effecting_op="false"} + live_done = f32[16] async-done(live_start), calls=async_comp + ROOT root = f32[8] copy(p0) +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloString)); + // `live_done` must check its `async-start` operand's attribute so `live_done` + // is not removed and `live_start` is not orphaned. + HloDCE dce(/*remove_cross_partition_collective_ops=*/true); + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_FALSE(changed); + EXPECT_NE(FindInstruction(module.get(), "live_done"), nullptr); + EXPECT_NE(FindInstruction(module.get(), "live_start"), nullptr); +} + +TEST_F(HloDceTest, DceAsyncWrappedCollective) { + constexpr absl::string_view kUnconstrainedHlo = R"hlo( +HloModule test_async_wrapped_collective_unconstrained + +async_comp { + p0 = f32[8] parameter(0) + ROOT ag = f32[16] all-gather(p0), dimensions={0}, replica_groups={{0,1}}, channel_id=1 +} + +ENTRY main { + p0 = f32[8] parameter(0) + dead_start = ((f32[8]), f32[16]) async-start(p0), calls=async_comp + dead_done = f32[16] async-done(dead_start), calls=async_comp + ROOT root = f32[8] copy(p0) +} +)hlo"; + // Preserved when remove_cross_partition_collective_ops=false. + { + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kUnconstrainedHlo)); + HloDCE dce(/*remove_cross_partition_collective_ops=*/false); + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_FALSE(changed); + EXPECT_NE(FindInstruction(module.get(), "dead_done"), nullptr); + EXPECT_NE(FindInstruction(module.get(), "dead_start"), nullptr); + } + + // Preserved when constrain_layout=true on the wrapped collective. + constexpr absl::string_view kConstrainedHlo = R"hlo( +HloModule test_async_wrapped_collective_constrained + +async_comp { + p0 = f32[8] parameter(0) + ROOT ag = f32[16] all-gather(p0), dimensions={0}, replica_groups={{0,1}}, channel_id=1, constrain_layout=true +} + +ENTRY main { + p0 = f32[8] parameter(0) + dead_start = ((f32[8]), f32[16]) async-start(p0), calls=async_comp + dead_done = f32[16] async-done(dead_start), calls=async_comp + ROOT root = f32[8] copy(p0) +} +)hlo"; + { + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kConstrainedHlo)); + HloDCE dce(/*remove_cross_partition_collective_ops=*/true); + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_FALSE(changed); + EXPECT_NE(FindInstruction(module.get(), "dead_done"), nullptr); + EXPECT_NE(FindInstruction(module.get(), "dead_start"), nullptr); + } + + // Eliminated when unconstrained and + // remove_cross_partition_collective_ops=true. + { + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kUnconstrainedHlo)); + HloDCE dce(/*remove_cross_partition_collective_ops=*/true); + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_TRUE(changed); + EXPECT_EQ(FindInstruction(module.get(), "dead_done"), nullptr); + EXPECT_EQ(FindInstruction(module.get(), "dead_start"), nullptr); + } +} + +TEST_F(HloDceTest, DceSelfAttributeSideEffectingCustomCallAndAsync) { + constexpr absl::string_view kHloString = R"hlo( +HloModule test_self_attr_true_custom_call_and_async + +async_comp { + p0 = f32[8] parameter(0) + ROOT inner_cc = f32[8] custom-call(p0), custom_call_target="tpu_custom_call", custom_call_has_side_effect=true +} + +ENTRY main { + p0 = f32[8] parameter(0) + dead_cc = f32[8] custom-call(p0), custom_call_target="tpu_custom_call", custom_call_has_side_effect=true, frontend_attributes={xla_allow_dce_side_effecting_op="true"} + dead_start = ((f32[8]), f32[8]) async-start(p0), calls=async_comp, frontend_attributes={xla_allow_dce_side_effecting_op="true"} + dead_done = f32[8] async-done(dead_start), calls=async_comp, frontend_attributes={xla_allow_dce_side_effecting_op="true"} + ROOT root = f32[8] copy(p0) +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloString)); + HloDCE dce; + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_TRUE(changed); + EXPECT_EQ(FindInstruction(module.get(), "dead_cc"), nullptr); + EXPECT_EQ(FindInstruction(module.get(), "dead_done"), nullptr); + EXPECT_EQ(FindInstruction(module.get(), "dead_start"), nullptr); +} + +TEST_F(HloDceTest, DcePairedAsyncDoneInheritingStartAttributeTrue) { + constexpr absl::string_view kHloString = R"hlo( +HloModule test_paired_async_done_inheriting_start_attr_true + +async_comp { + p0 = f32[8] parameter(0) + ROOT inner_cc = f32[8] custom-call(p0), custom_call_target="tpu_custom_call", custom_call_has_side_effect=true +} + +ENTRY main { + p0 = f32[8] parameter(0) + dead_start = ((f32[8]), f32[8]) async-start(p0), calls=async_comp, frontend_attributes={xla_allow_dce_side_effecting_op="true"} + dead_done = f32[8] async-done(dead_start), calls=async_comp + ROOT root = f32[8] copy(p0) +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloString)); + HloDCE dce; + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_TRUE(changed); + EXPECT_EQ(FindInstruction(module.get(), "dead_done"), nullptr); + EXPECT_EQ(FindInstruction(module.get(), "dead_start"), nullptr); +} + +TEST_F(HloDceTest, DceSinglePassReversePostOrderSideEffectingChain) { + constexpr absl::string_view kHloString = R"hlo( +HloModule test_single_pass_reverse_postorder_chain + +async_comp { + p0 = f32[8] parameter(0) + ROOT inner_cc = f32[8] custom-call(p0), custom_call_target="tpu_custom_call", custom_call_has_side_effect=true +} + +ENTRY main { + p0 = f32[8] parameter(0) + producer_cc = f32[8] custom-call(p0), custom_call_target="tpu_custom_call", custom_call_has_side_effect=true, frontend_attributes={xla_allow_dce_side_effecting_op="true"} + dead_start = ((f32[8]), f32[8]) async-start(producer_cc), calls=async_comp, frontend_attributes={xla_allow_dce_side_effecting_op="true"} + dead_done = f32[8] async-done(dead_start), calls=async_comp, frontend_attributes={xla_allow_dce_side_effecting_op="true"} + ROOT root = f32[8] copy(p0) +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloString)); + HloDCE dce; + // Reverse postorder visits consumers before operands (`dead_done` -> + // `dead_start` -> `producer_cc`), allowing a single pass to eliminate the + // entire chain even though `RemoveInstructionAndUnusedOperands` does not + // recurse into side-effecting operands. + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_TRUE(changed); + EXPECT_EQ(FindInstruction(module.get(), "dead_done"), nullptr); + EXPECT_EQ(FindInstruction(module.get(), "dead_start"), nullptr); + EXPECT_EQ(FindInstruction(module.get(), "producer_cc"), nullptr); +} + } // namespace } // namespace xla From 71b4cb9b05d8a982c57ccfcd65e8b4fb264e8c6f Mon Sep 17 00:00:00 2001 From: Junwhan Ahn Date: Wed, 23 Sep 2026 11:17:52 -0700 Subject: [PATCH 26/28] [IFRT] Refactor tf2hlo_test to use real PJRT-CPU IFRT client PiperOrigin-RevId: 986883135 --- tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD | 3 +-- .../compiler/mlir/tfrt/transforms/ifrt/tf2hlo_test.cc | 10 +++------- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD index 7ae680f60668c2..3fc654d38b506d 100644 --- a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD +++ b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD @@ -241,10 +241,9 @@ tf_cc_test( "@xla//xla/pjrt:pjrt_compiler", "@xla//xla/pjrt/plugin/xla_cpu:cpu_topology_description", "@xla//xla/python/ifrt", - "@xla//xla/python/ifrt:mock", "@xla//xla/python/ifrt:test_util", "@xla//xla/python/pjrt_ifrt", - "@xla//xla/python/pjrt_ifrt:tfrt_cpu_client_test_lib", + "@xla//xla/python/pjrt_ifrt:pjrt_cpu_client_test_lib", "@xla//xla/service:device_assignment", "@xla//xla/tsl/platform:statusor", ], diff --git a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo_test.cc b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo_test.cc index d5ac5184661058..780e361f06f2dd 100644 --- a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo_test.cc +++ b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo_test.cc @@ -43,7 +43,6 @@ limitations under the License. #include "xla/pjrt/pjrt_compiler.h" #include "xla/pjrt/plugin/xla_cpu/cpu_topology_description.h" #include "xla/python/ifrt/client.h" -#include "xla/python/ifrt/mock.h" #include "xla/python/ifrt/test_util.h" #include "xla/python/pjrt_ifrt/pjrt_topology.h" #include "xla/service/device_assignment.h" @@ -514,18 +513,15 @@ TEST_F(Tf2HloTest, GpuCompile) { ASSERT_TRUE(mlir_module); ASSERT_TRUE(mlir_module.get() != nullptr); - xla::ifrt::MockClient mock_client; - ON_CALL(mock_client, GetDefaultDeviceAssignment) - .WillByDefault([]() -> absl::StatusOr { - return xla::DeviceAssignment(1, 1); - }); + TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr client, + xla::ifrt::test_util::GetClient()); std::vector dtype_and_shapes; dtype_and_shapes.push_back(DtypeAndShape{DT_FLOAT, {}}); TF_ASSERT_OK_AND_ASSIGN( tensorflow::tpu::TPUCompileMetadataProto compile_metadata, - GetCompileMetadata(mlir_module.get(), mock_client)); + GetCompileMetadata(mlir_module.get(), *client)); TF_ASSERT_OK(UpdateCompileMetadata(compile_metadata, dtype_and_shapes)); std::vector variable_arg_indices; From 0374be1c844ce7be1bb986191df805f85c78fc4f Mon Sep 17 00:00:00 2001 From: Bixia Zheng Date: Wed, 23 Sep 2026 11:18:43 -0700 Subject: [PATCH 27/28] Clear sharding before spmd-cleanup. After GSPMD generates device local, it invokes passes such as DCE to clean up the code and then clears sharding. Previously, it first cleans up the code then clears sharding. This CL changes GSPMD to first clear sharding then clean up the code, so that the clean up passes will not see sharding as the code is already in device local representation. PiperOrigin-RevId: 986883751 --- .../xla/xla/service/spmd/spmd_partitioner.cc | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.cc b/third_party/xla/xla/service/spmd/spmd_partitioner.cc index 2b6feb4ca0b595..a866832ef944b1 100644 --- a/third_party/xla/xla/service/spmd/spmd_partitioner.cc +++ b/third_party/xla/xla/service/spmd/spmd_partitioner.cc @@ -7405,15 +7405,6 @@ absl::StatusOr SpmdPartitioner::RunImpl( *module, options_.report_instruction_count)); XLA_VLOG_LINES(1, logger.MakeReport()); - if (changed) { - HloPassPipeline pass("spmd-cleanup"); - pass.AddPass(/*remove_cross_partition_collective_ops=*/true); - pass.AddPass(); - pass.AddPass(/*remove_cross_partition_collective_ops=*/true); - pass.AddPass(/*is_layout_sensitive=*/false); - ABSL_RETURN_IF_ERROR(pass.Run(module, execution_threads).status()); - } - ABSL_RETURN_IF_ERROR(ClearShardingAttributes( module, num_replicas() * num_partitions(), execution_threads)); @@ -7424,6 +7415,15 @@ absl::StatusOr SpmdPartitioner::RunImpl( entry_root->set_sharding(std::move(final_sharding)); } + if (changed) { + HloPassPipeline pass("spmd-cleanup"); + pass.AddPass(/*remove_cross_partition_collective_ops=*/true); + pass.AddPass(); + pass.AddPass(/*remove_cross_partition_collective_ops=*/true); + pass.AddPass(/*is_layout_sensitive=*/false); + ABSL_RETURN_IF_ERROR(pass.Run(module, execution_threads).status()); + } + return changed; } From 84d17f0377a8efce6ec561a05c81fd64dc77f2ed Mon Sep 17 00:00:00 2001 From: Viktor Moros Date: Wed, 23 Sep 2026 12:05:02 -0700 Subject: [PATCH 28/28] [XLA Highway] Add 256-bit rectangular transpose kernel PiperOrigin-RevId: 986914259 --- third_party/xla/xla/pjrt/transpose_kernels.h | 71 ++++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/third_party/xla/xla/pjrt/transpose_kernels.h b/third_party/xla/xla/pjrt/transpose_kernels.h index 59f6b37361b5b2..90212a50fda7a4 100644 --- a/third_party/xla/xla/pjrt/transpose_kernels.h +++ b/third_party/xla/xla/pjrt/transpose_kernels.h @@ -333,6 +333,77 @@ HWY_FLATTEN void TransposeMicroKernel256Square(const char* HWY_RESTRICT a, }; BuildVecs256SquareCPS(d, d_half, a, lda, cont); } +// Permutes 64-bit lanes across 128-bit blocks in a 256-bit vector. +// Specifically, reorders the 4 64-bit lanes from [0, 1, 2, 3] to [0, 2, 1, 3]. +template +HWY_INLINE V ShuffleAcrossBlocks_3120(V v) { + const hn::DFromV d; + const hn::Repartition d64; + return hn::BitCast(d, + hn::Per4LaneBlockShuffle<3, 1, 2, 0>(hn::BitCast(d64, v))); +} + +template +HWY_INLINE void StoreRectPairs(DHalf d_half, char* b, int64_t ldb, Vs... v) { + ForEachIndexed( + [&](auto idx_c, auto&& vec) { + constexpr size_t i = decltype(idx_c)::value; + const auto lo = hn::LowerHalf(d_half, vec); + const auto hi = hn::UpperHalf(d_half, vec); + hn::StoreU(lo, d_half, + reinterpret_cast*>(b + ldb * (i * 2))); + hn::StoreU(hi, d_half, + reinterpret_cast*>(b + ldb * (i * 2 + 1))); + }, + v...); +} + +// Loads a rectangular block of memory into 256-bit vectors. Invokes the +// continuation with the constructed sequence of vectors. +template +HWY_INLINE void BuildVecs256RectCPS(D d, DHalf d_half, const char* a, + int64_t lda, Cont& cont, Vs... vs) { + if constexpr (I == bs / 2) { + cont(vs...); + } else { + auto lo = hn::LoadU( + d_half, reinterpret_cast*>(a + lda * I)); + auto hi = hn::LoadU(d_half, reinterpret_cast*>( + a + lda * (I + bs / 2))); + auto v = hn::Combine(d, hi, lo); + BuildVecs256RectCPS(d, d_half, a, lda, cont, vs..., v); + } +} + +template +HWY_INLINE void BuildVecs256RectCPS(D d, DHalf d_half, const char* a, + int64_t lda, Cont& cont) { + BuildVecs256RectCPS<0, T, bs>(d, d_half, a, lda, cont); +} + +template +HWY_FLATTEN void TransposeMicroKernel256Rect(const char* HWY_RESTRICT a, + int64_t lda, char* HWY_RESTRICT b, + int64_t ldb) { + using D = hn::FixedTag; + using DHalf = hn::Half; + static constexpr size_t element_size = sizeof(T); + static constexpr size_t row_bytes = element_size * bs; + static_assert(row_bytes == kBlockBytes); + static_assert(bs % 2 == 0); + const D d; + const DHalf d_half; + auto cont = [&](auto... vecs) { + InterleaveWithinBlocksCPS( + [&](auto... interleaved) { + StoreRectPairs(d_half, b, ldb, + ShuffleAcrossBlocks_3120(interleaved)...); + }, + vecs...); + }; + BuildVecs256RectCPS(d, d_half, a, lda, cont); +} #endif // HWY_MIN_BYTES >= 32 } // namespace HWY_NAMESPACE