From 15a87d5e28b689c7cc5462ea804a39e7ba4a1df4 Mon Sep 17 00:00:00 2001 From: Maddipatla Chatan Date: Sat, 19 Sep 2026 21:49:24 +0530 Subject: [PATCH 01/17] Fix heartbeat.py logging and return statement --- tensorflow/dtensor/python/heartbeat.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow/dtensor/python/heartbeat.py b/tensorflow/dtensor/python/heartbeat.py index d3c8c3a247b15f..d1f74c933f8f0c 100644 --- a/tensorflow/dtensor/python/heartbeat.py +++ b/tensorflow/dtensor/python/heartbeat.py @@ -85,6 +85,7 @@ def _heartbeat( logging.warning('Heartbeat failure %d, %d more until limit: %s', _failure_count, _CONSECUTIVE_FAILURES_LIMIT - _failure_count, e) + continue else: logging.fatal('Heartbeat failure %d, limit of %d reached: %s', _failure_count, _CONSECUTIVE_FAILURES_LIMIT, e) From 6459fce91044a5d70707286dcc2138ba015c73f7 Mon Sep 17 00:00:00 2001 From: Henning Becker Date: Mon, 21 Sep 2026 05:18:51 -0700 Subject: [PATCH 02/17] Rename `KernelArguments::OperandIndex`/`ResultIndex` to `PositionOfOperand`/`PositionOfResult`. These accessors map an HLO operand number, or a result shape index, to a position in the kernel argument list. The value going in and the value coming out live in different index spaces, and calling both of them an "index" invited the reader to mix them up. `KernelArguments` already calls the outgoing value a position internally (`ArgumentPositions`), so this brings the public names in line with that vocabulary. Pure rename, no behavior change. PiperOrigin-RevId: 985207402 --- .../xla/codegen/emitters/kernel_arguments.cc | 15 ++++---- .../xla/codegen/emitters/kernel_arguments.h | 17 ++++++---- .../codegen/emitters/kernel_arguments_test.cc | 34 +++++++++---------- 3 files changed, 35 insertions(+), 31 deletions(-) diff --git a/third_party/xla/xla/codegen/emitters/kernel_arguments.cc b/third_party/xla/xla/codegen/emitters/kernel_arguments.cc index 41d946c8b7b25e..2f040807966386 100644 --- a/third_party/xla/xla/codegen/emitters/kernel_arguments.cc +++ b/third_party/xla/xla/codegen/emitters/kernel_arguments.cc @@ -307,22 +307,23 @@ absl::StatusOr KernelArguments::Create( interleaved_output_indices); } -absl::StatusOr KernelArguments::OperandIndex( - int64_t operand_index) const { +absl::StatusOr KernelArguments::PositionOfOperand( + int64_t operand_number_in_hlo) const { if (!positions_.has_value()) { return absl::FailedPreconditionError( "KernelArguments were not created from an HLO instruction, so operand " "positions are unknown"); } - if (operand_index < 0 || - operand_index >= static_cast(positions_->operands.size())) { + if (operand_number_in_hlo < 0 || + operand_number_in_hlo >= + static_cast(positions_->operands.size())) { return absl::OutOfRangeError( - absl::StrCat("No such operand: ", operand_index)); + absl::StrCat("No such operand: ", operand_number_in_hlo)); } - return positions_->operands[operand_index]; + return positions_->operands[operand_number_in_hlo]; } -absl::StatusOr KernelArguments::ResultIndex( +absl::StatusOr KernelArguments::PositionOfResult( const ShapeIndex& shape_index) const { if (!positions_.has_value()) { return absl::FailedPreconditionError( diff --git a/third_party/xla/xla/codegen/emitters/kernel_arguments.h b/third_party/xla/xla/codegen/emitters/kernel_arguments.h index f8b131160854e2..4de8500c6022b0 100644 --- a/third_party/xla/xla/codegen/emitters/kernel_arguments.h +++ b/third_party/xla/xla/codegen/emitters/kernel_arguments.h @@ -162,8 +162,9 @@ class KernelArguments { const std::vector& args() const { return args_; } - // Position of the argument that carries operand `operand_index` of the - // instruction these arguments were created for. + // Position, in this argument list, of the argument that carries operand + // `operand_number_in_hlo` of the HLO instruction these arguments were + // created for. // // Use this instead of hardcoding argument positions when building a // `stream_executor::KernelArgsPackingSpec`, whose relocation indices refer @@ -171,15 +172,17 @@ class KernelArguments { // // Returns `FailedPrecondition` if the arguments were not built by one of the // `Create` factories, and `OutOfRange` if there is no such operand. - absl::StatusOr OperandIndex(int64_t operand_index) const; + absl::StatusOr PositionOfOperand( + int64_t operand_number_in_hlo) const; - // Position of the argument that carries the result subshape at - // `shape_index`. For an instruction with a non-tuple result, that is - // `ResultIndex({})`; for a one-element tuple result, `ResultIndex({0})`. + // Position, in this argument list, of the argument that carries the result + // subshape at `shape_index` of the HLO instruction these arguments were + // created for. For a non-tuple result that is `PositionOfResult({})`; for a + // one-element tuple result, `PositionOfResult({0})`. // // Returns `FailedPrecondition` if the arguments were not built by one of the // `Create` factories, and `OutOfRange` if there is no such result. - absl::StatusOr ResultIndex(const ShapeIndex& shape_index) const; + absl::StatusOr PositionOfResult(const ShapeIndex& shape_index) const; std::vector GetArgumentShapedSlices() const { std::vector arg_slices; diff --git a/third_party/xla/xla/codegen/emitters/kernel_arguments_test.cc b/third_party/xla/xla/codegen/emitters/kernel_arguments_test.cc index 423fe75fd1bebc..952d54d7923fe5 100644 --- a/third_party/xla/xla/codegen/emitters/kernel_arguments_test.cc +++ b/third_party/xla/xla/codegen/emitters/kernel_arguments_test.cc @@ -589,21 +589,21 @@ TEST_F(KernelArgumentsTest, OperandAndResultIndices) { root)); // Operands come first, then the array leaves of the result shape. - EXPECT_THAT(kernel_arguments.OperandIndex(0), IsOkAndHolds(0)); - EXPECT_THAT(kernel_arguments.OperandIndex(1), IsOkAndHolds(1)); - EXPECT_THAT(kernel_arguments.ResultIndex({0}), IsOkAndHolds(2)); - EXPECT_THAT(kernel_arguments.ResultIndex({1}), IsOkAndHolds(3)); + EXPECT_THAT(kernel_arguments.PositionOfOperand(0), IsOkAndHolds(0)); + EXPECT_THAT(kernel_arguments.PositionOfOperand(1), IsOkAndHolds(1)); + EXPECT_THAT(kernel_arguments.PositionOfResult({0}), IsOkAndHolds(2)); + EXPECT_THAT(kernel_arguments.PositionOfResult({1}), IsOkAndHolds(3)); - EXPECT_THAT(kernel_arguments.OperandIndex(2), + EXPECT_THAT(kernel_arguments.PositionOfOperand(2), StatusIs(absl::StatusCode::kOutOfRange)); - EXPECT_THAT(kernel_arguments.OperandIndex(-1), + EXPECT_THAT(kernel_arguments.PositionOfOperand(-1), StatusIs(absl::StatusCode::kOutOfRange)); // The root is a tuple, so there is no array leaf at the top level. - EXPECT_THAT(kernel_arguments.ResultIndex({}), + EXPECT_THAT(kernel_arguments.PositionOfResult({}), StatusIs(absl::StatusCode::kOutOfRange)); } -TEST_F(KernelArgumentsTest, OperandAndResultIndicesFollowInterleaving) { +TEST_F(KernelArgumentsTest, OperandAndResultPositionsFollowInterleaving) { ASSERT_OK_AND_ASSIGN(std::unique_ptr module, ParseAndReturnVerifiedModule(R"( ENTRY main { @@ -629,13 +629,13 @@ TEST_F(KernelArgumentsTest, OperandAndResultIndicesFollowInterleaving) { KernelArguments::Create(*assignment, gpu::GetDefaultBufferAlignment(), root, interleaved_indices)); - EXPECT_THAT(kernel_arguments.OperandIndex(0), IsOkAndHolds(0)); - EXPECT_THAT(kernel_arguments.ResultIndex({0}), IsOkAndHolds(1)); - EXPECT_THAT(kernel_arguments.OperandIndex(1), IsOkAndHolds(2)); - EXPECT_THAT(kernel_arguments.ResultIndex({1}), IsOkAndHolds(3)); + EXPECT_THAT(kernel_arguments.PositionOfOperand(0), IsOkAndHolds(0)); + EXPECT_THAT(kernel_arguments.PositionOfResult({0}), IsOkAndHolds(1)); + EXPECT_THAT(kernel_arguments.PositionOfOperand(1), IsOkAndHolds(2)); + EXPECT_THAT(kernel_arguments.PositionOfResult({1}), IsOkAndHolds(3)); } -TEST_F(KernelArgumentsTest, ResultIndexOfNonTupleResult) { +TEST_F(KernelArgumentsTest, PositionOfResultOfNonTupleResult) { ASSERT_OK_AND_ASSIGN(std::unique_ptr module, ParseAndReturnVerifiedModule(R"( ENTRY main { @@ -658,8 +658,8 @@ TEST_F(KernelArgumentsTest, ResultIndexOfNonTupleResult) { KernelArguments::Create(*assignment, gpu::GetDefaultBufferAlignment(), root)); - EXPECT_THAT(kernel_arguments.ResultIndex({}), IsOkAndHolds(2)); - EXPECT_THAT(kernel_arguments.ResultIndex({0}), + EXPECT_THAT(kernel_arguments.PositionOfResult({}), IsOkAndHolds(2)); + EXPECT_THAT(kernel_arguments.PositionOfResult({0}), StatusIs(absl::StatusCode::kOutOfRange)); } @@ -667,9 +667,9 @@ TEST_F(KernelArgumentsTest, HandBuiltArgumentsHaveNoKnownPositions) { KernelArguments kernel_arguments(std::vector{KernelArgument( ShapeUtil::MakeShape(F32, {4}), BufferAllocation::Slice())}); - EXPECT_THAT(kernel_arguments.OperandIndex(0), + EXPECT_THAT(kernel_arguments.PositionOfOperand(0), StatusIs(absl::StatusCode::kFailedPrecondition)); - EXPECT_THAT(kernel_arguments.ResultIndex({}), + EXPECT_THAT(kernel_arguments.PositionOfResult({}), StatusIs(absl::StatusCode::kFailedPrecondition)); } From 5ee725504580b191430b79743dc5f2be4ef923ac Mon Sep 17 00:00:00 2001 From: Stanislav Bardyuk Date: Mon, 21 Sep 2026 05:37:56 -0700 Subject: [PATCH 03/17] PR #47282: [XLA:TSL] Fix silent truncation of ZlibOutputBuffer appends larger than 4 GiB 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/47282 ## πŸ“ Summary of Changes `ZlibOutputBuffer::Append` passes data that does not fit in the staging buffer directly to zlib by assigning the `size_t` byte count to `z_stream::avail_in`, which is zlib's 32-bit `uInt`. A single `Append` of 4 GiB or more is truncated to `size % 2^32` bytes: the write reports OK, `Close()` produces a well-formed gzip stream, and everything past the truncated count is silently lost. This change chunks the direct-deflate path so zlib is fed at most `uInt` max bytes per round, and adds a round-trip regression test that appends 4 GiB + 1 MiB in one call. The test lives in its own `size = "large"` target (`zlib_buffers_large_append_test`) because it needs ~4 GiB of memory for the input string; the existing `zlib_buffers_test` stays `small`. This completes the fix started in 783463ddeb, which addressed the size comparisons but not the `avail_in` assignment (see the reporter's follow-up in the issue). ## 🎯 Justification Fixes #28159: profiling a large workload makes `jax.profiler.stop_trace()` write the whole Chrome-trace JSON through one `Append` call in `WriteGzippedDataToFile` (`xla/tsl/profiler/rpc/client/save_profile.cc`); above 4 GiB the resulting `.trace.json.gz` decompresses to a truncated JSON document. Fixing `Append` itself (rather than chunking in `save_profile.cc`) also covers the other users of this class (`record_writer`, `xla/service/dump.cc`). Reproduced at HEAD: appending 2^32 + 2^20 bytes produced a file that decompresses to exactly 2^20 bytes (`size % 2^32`). ## πŸš€ Kind of Contribution πŸ› Bug Fix ## πŸ§ͺ Unit Tests: - New target `//xla/tsl/lib/io:zlib_buffers_large_append_test` (`size = "large"`, ~4 GiB peak RSS, 8 s at `Z_BEST_SPEED`): single 4 GiB + 1 MiB `Append`, round-tripped through `ZlibInputStream` with full content verification and an end-of-stream check. Fails before this change (first read hits `OUT_OF_RANGE: EOF reached` - the file decompresses to exactly `size % 2^32` = 1 MiB), passes after. - All 13 test targets in `//xla/tsl/lib/io/...` pass with the change. - For appends below 4 GiB the loop runs exactly one iteration with the same zlib call sequence as before, so existing behavior is unchanged (covered by the existing round-trip tests, which exercise the direct-deflate branch with data larger than the staging buffer). ## πŸ§ͺ Execution Tests: Not applicable (host-side I/O library; no device execution involved). Copybara import of the project: -- 391f419e530e1ed5184d114f29eadc2354abd631 by Stanislav Bardyuk : Fix silent truncation of ZlibOutputBuffer appends larger than 4 GiB ZlibOutputBuffer::Append assigns the size_t byte count of data that does not fit in the staging buffer directly to zlib's 32-bit avail_in, so a single Append of 4 GiB or more is truncated to size % 2^32 bytes while reporting success and producing a well-formed gzip stream. This is how jax.profiler.stop_trace() corrupts .trace.json.gz for large traces: the whole trace JSON goes through one Append in WriteGzippedDataToFile. Feed the direct-deflate path at most UINT_MAX bytes per round instead, and add a regression test that round-trips a single 4 GiB + 1 MiB append (own size = "large" target; the test needs ~4 GiB for the input string). Completes 783463ddeb, which fixed the size comparisons but not the avail_in assignment. Fixes #28159. Merging this change closes #47282 PiperOrigin-RevId: 985215524 --- third_party/xla/xla/tsl/lib/io/BUILD | 24 ++++- .../lib/io/zlib_buffers_large_append_test.cc | 101 ++++++++++++++++++ .../xla/xla/tsl/lib/io/zlib_outputbuffer.cc | 38 ++++--- 3 files changed, 149 insertions(+), 14 deletions(-) create mode 100644 third_party/xla/xla/tsl/lib/io/zlib_buffers_large_append_test.cc diff --git a/third_party/xla/xla/tsl/lib/io/BUILD b/third_party/xla/xla/tsl/lib/io/BUILD index 209a2a6810e9ef..210dd2e16d9983 100644 --- a/third_party/xla/xla/tsl/lib/io/BUILD +++ b/third_party/xla/xla/tsl/lib/io/BUILD @@ -14,7 +14,7 @@ # ============================================================================== load("@xla//third_party/rules_python/python:py_test.bzl", "py_test") -load("//xla/tsl:tsl.bzl", "internal_visibility") +load("//xla/tsl:tsl.bzl", "if_google", "internal_visibility") load("//xla/tsl:tsl.default.bzl", "filegroup") load("//xla/tsl/platform:build_config.bzl", "tsl_cc_test") load( @@ -654,6 +654,28 @@ tsl_cc_test( ], ) +tsl_cc_test( + name = "zlib_buffers_large_append_test", + size = "large", + srcs = ["zlib_buffers_large_append_test.cc"], + tags = if_google([ + "requires-mem:16g", + ]), + deps = [ + ":random_inputstream", + ":zlib_compression_options", + ":zlib_inputstream", + ":zlib_outputbuffer", + "//xla/tsl/platform:env", + "//xla/tsl/platform:test", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", + "@com_google_googletest//:gtest_main", + "@tsl//tsl/platform:tstring", + "@zlib", + ], +) + cc_library( name = "zip_writer", srcs = ["zip_writer.cc"], diff --git a/third_party/xla/xla/tsl/lib/io/zlib_buffers_large_append_test.cc b/third_party/xla/xla/tsl/lib/io/zlib_buffers_large_append_test.cc new file mode 100644 index 00000000000000..993690db77d3c6 --- /dev/null +++ b/third_party/xla/xla/tsl/lib/io/zlib_buffers_large_append_test.cc @@ -0,0 +1,101 @@ +/* Copyright 2026 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include + +#include +#include +#include +#include + +#include "absl/status/status.h" +#include "absl/strings/string_view.h" +#include "xla/tsl/lib/io/random_inputstream.h" +#include "xla/tsl/lib/io/zlib_compression_options.h" +#include "xla/tsl/lib/io/zlib_inputstream.h" +#include "xla/tsl/lib/io/zlib_outputbuffer.h" +#include "xla/tsl/platform/env.h" +#include "xla/tsl/platform/file_system.h" +#include "xla/tsl/platform/test.h" +#include "tsl/platform/tstring.h" + +namespace tsl { +namespace io { +namespace { + +// A single Append larger than 4 GiB used to be silently truncated to +// data.size() % 4 GiB bytes because zlib's avail_in is 32-bit. +// This test needs ~4.5 GiB of memory for the input string. +TEST(ZlibBuffers, AppendLargerThan4GB) { + Env* env = Env::Default(); + std::string fname; + ASSERT_TRUE(env->LocalTempFilename(&fname)); + + constexpr size_t kDataSize = (1ULL << 32) + (1ULL << 20); + constexpr absl::string_view kPattern = + "The quick brown fox jumps over the lazy dog. "; + std::string data; + data.reserve(kDataSize); + data.append(kPattern); + while (data.size() < kDataSize) { + data.append(data.data(), std::min(data.size(), kDataSize - data.size())); + } + + ZlibCompressionOptions output_options = ZlibCompressionOptions::GZIP(); + output_options.compression_level = Z_BEST_SPEED; + + std::unique_ptr file_writer; + ASSERT_OK(env->NewWritableFile(fname, &file_writer)); + ZlibOutputBuffer out(file_writer.get(), output_options.input_buffer_size, + output_options.output_buffer_size, output_options); + ASSERT_OK(out.Init()); + ASSERT_OK(out.Append(absl::string_view(data))); + ASSERT_OK(out.Close()); + ASSERT_OK(file_writer->Flush()); + ASSERT_OK(file_writer->Close()); + + ZlibCompressionOptions input_options = ZlibCompressionOptions::GZIP(); + std::unique_ptr file_reader; + ASSERT_OK(env->NewRandomAccessFile(fname, &file_reader)); + std::unique_ptr input_stream( + new RandomAccessInputStream(file_reader.get())); + ZlibInputStream in(input_stream.get(), input_options.input_buffer_size, + input_options.output_buffer_size, input_options); + + // Read back in chunks so peak memory stays close to the input string. + constexpr size_t kChunkSize = 1 << 24; + size_t total_read = 0; + tstring chunk; + while (total_read < kDataSize) { + size_t bytes_to_read = std::min(kChunkSize, kDataSize - total_read); + absl::Status read_status = in.ReadNBytes(bytes_to_read, &chunk); + ASSERT_TRUE(read_status.ok()) + << "read failed after " << total_read << " bytes: " << read_status; + ASSERT_EQ(chunk.size(), bytes_to_read); + // Compare without ASSERT_EQ on the strings to avoid logging the 16 MiB + // operands on mismatch. + ASSERT_TRUE(absl::string_view(chunk) == + absl::string_view(data).substr(total_read, bytes_to_read)) + << "content mismatch at offset " << total_read; + total_read += bytes_to_read; + } + // The decompressed stream must end exactly at kDataSize bytes. + EXPECT_TRUE(absl::IsOutOfRange(in.ReadNBytes(1, &chunk))); + EXPECT_OK(env->DeleteFile(fname)); +} + +} // namespace +} // namespace io +} // namespace tsl diff --git a/third_party/xla/xla/tsl/lib/io/zlib_outputbuffer.cc b/third_party/xla/xla/tsl/lib/io/zlib_outputbuffer.cc index c873ead32b7ce9..d87697a7dfd41e 100644 --- a/third_party/xla/xla/tsl/lib/io/zlib_outputbuffer.cc +++ b/third_party/xla/xla/tsl/lib/io/zlib_outputbuffer.cc @@ -17,9 +17,11 @@ limitations under the License. #include +#include #include #include #include +#include #include #include "absl/log/check.h" @@ -191,19 +193,29 @@ absl::Status ZlibOutputBuffer::Append(absl::string_view data) { // `data` is too large to fit in input buffer so we deflate it directly. // Note that at this point we have already deflated all existing input so // we do not need to backup next_in and avail_in. - z_stream_->next_in = reinterpret_cast(const_cast(data.data())); - z_stream_->avail_in = bytes_to_write; - - do { - if (z_stream_->avail_out == 0) { - // No available output space. - // Write output buffer to file. - ABSL_RETURN_IF_ERROR(FlushOutputBufferToFile()); - } - ABSL_RETURN_IF_ERROR(Deflate(zlib_options_.flush_mode)); - } while (z_stream_->avail_out == 0); - - DCHECK(z_stream_->avail_in == 0); // All input will be used up. + // zlib's avail_in is a uInt, so feed it at most UINT_MAX bytes at a time + // to avoid truncating the byte count for larger appends. + size_t bytes_deflated = 0; + while (bytes_deflated < bytes_to_write) { + // NOLINTNEXTLINE(misc-include-cleaner): uInt comes from . + const uInt chunk = static_cast(std::min( + bytes_to_write - bytes_deflated, std::numeric_limits::max())); + z_stream_->next_in = reinterpret_cast( + const_cast(data.data() + bytes_deflated)); + z_stream_->avail_in = chunk; + + do { + if (z_stream_->avail_out == 0) { + // No available output space. + // Write output buffer to file. + ABSL_RETURN_IF_ERROR(FlushOutputBufferToFile()); + } + ABSL_RETURN_IF_ERROR(Deflate(zlib_options_.flush_mode)); + } while (z_stream_->avail_out == 0); + + DCHECK(z_stream_->avail_in == 0); // All input will be used up. + bytes_deflated += chunk; + } // Restore z_stream input pointers. z_stream_->next_in = z_stream_input_.get(); From bf7dfc3f10247a1a1533e08bd9430cfc14920866 Mon Sep 17 00:00:00 2001 From: spesherov <9691985+spesherov@users.noreply.github.com> Date: Mon, 21 Sep 2026 05:57:33 -0700 Subject: [PATCH 04/17] PR #47219: Upgrade curl to 8.21.0 and update libcurl integration. 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/47219 πŸ“ Summary of Changes Upgrade curl from the previous pin to **8.21.0** and adapt XLA/TSL libcurl integration. - Bump curl archive pin in `workspace2.bzl` (`strip_prefix`, URL, SHA256). - Refresh `third_party/curl.BUILD` for the 8.21.0 source layout (explicit `lib/`/`src/` lists; vquic `cf-*` renames and new modules). - Update `CurlHttpRequest` for newer libcurl APIs (`CURLOPT_UPLOAD` instead of deprecated `CURLOPT_PUT`; option/signature type updates). - Adjust `curl_http_request_test` fakes for the updated options/signatures. - `USE_SECTRANSP` was Secure Transport, which curl 8.15 removed. `USE_APPLE_SECTRUST` is curl 8.17’s replacement for using Apple’s OS trust store for certificate verification. 🎯 Justification Keeps XLA/TSL’s vendored curl up to date for security and compatibility. Curl releases after the previous pin fix multiple CVEs (credential leaks, connection reuse, TLS/HTTP2 issues, etc.). TensorFlow imports XLA via Copybara, so these changes must land in openxla/xla first; companion TF PR: https://github.com/tensorflow/tensorflow/pull/115842 πŸš€ Kind of Contribution ♻️ Cleanup (dependency/security upgrade) πŸ“Š Benchmark (for Performance Improvements) N/A β€” dependency upgrade, not a performance change. πŸ§ͺ Unit Tests: No new XLA compiler unit tests. Existing `curl_http_request_test` updated for the new libcurl option types/signatures. πŸ§ͺ Execution Tests: N/A β€” no HLO/GPU execution changes. Validated locally by building `@curl` / curl HTTP request targets from the companion TensorFlow tree. Copybara import of the project: -- f2754cd3bdd4b1ed36c9dbb3ed555c2e1bf8eaec by spesherov <9691985+spesherov@users.noreply.github.com>: Upgrade curl to 8.21.0 and update libcurl integration. Bump the curl archive pin and refresh curl.BUILD for the 8.21.0 source layout. Update CurlHttpRequest for CURLOPT_UPLOAD and newer libcurl option/signature types. Merging this change closes #47219 PiperOrigin-RevId: 985224070 --- tensorflow/workspace2.bzl | 6 +- third_party/xla/third_party/curl.BUILD | 236 ++++++++++-------- third_party/xla/workspace2.bzl | 6 +- .../tsl/platform/cloud/curl_http_request.cc | 19 +- .../platform/cloud/curl_http_request_test.cc | 13 +- 5 files changed, 156 insertions(+), 124 deletions(-) diff --git a/tensorflow/workspace2.bzl b/tensorflow/workspace2.bzl index d2dad0607e04e6..8e878d43070531 100644 --- a/tensorflow/workspace2.bzl +++ b/tensorflow/workspace2.bzl @@ -423,10 +423,10 @@ def _tf_repositories(): tf_http_archive( name = "curl", build_file = "@xla//third_party:curl.BUILD", - sha256 = "264537d90e58d2b09dddc50944baf3c38e7089151c8986715e2aaeaaf2b8118f", - strip_prefix = "curl-8.11.0", + sha256 = "d9b327997999045a24cda50f3983e69e51c516bd8be6ef9842fc7f99135e33bb", + strip_prefix = "curl-8.21.0", system_build_file = "//third_party/systemlibs:curl.BUILD", - urls = tf_mirror_urls("https://curl.se/download/curl-8.11.0.tar.gz"), + urls = tf_mirror_urls("https://curl.se/download/curl-8.21.0.tar.gz"), ) tf_http_archive( diff --git a/third_party/xla/third_party/curl.BUILD b/third_party/xla/third_party/curl.BUILD index 0aaf949044a875..917bf2b755e172 100644 --- a/third_party/xla/third_party/curl.BUILD +++ b/third_party/xla/third_party/curl.BUILD @@ -35,16 +35,15 @@ CURL_WIN_COPTS = [ ] CURL_WIN_SRCS = [ - "lib/asyn-thread.c", - "lib/inet_ntop.c", + "lib/dllmain.c", "lib/system_win32.c", - "lib/setup-win32.h", ] cc_library( name = "curl", srcs = [ "include/curl_config.h", + ] + [ "lib/altsvc.c", "lib/altsvc.h", "lib/amigaos.c", @@ -52,13 +51,14 @@ cc_library( "lib/arpa_telnet.h", "lib/asyn.h", "lib/asyn-ares.c", - "lib/base64.c", + "lib/asyn-base.c", + "lib/asyn-thrdd.c", "lib/bufq.c", "lib/bufq.h", "lib/bufref.c", "lib/bufref.h", - "lib/c-hyper.c", - "lib/c-hyper.h", + "lib/cf-dns.c", + "lib/cf-dns.h", "lib/cf-h1-proxy.c", "lib/cf-h1-proxy.h", "lib/cf-h2-proxy.c", @@ -67,18 +67,20 @@ cc_library( "lib/cf-haproxy.h", "lib/cf-https-connect.c", "lib/cf-https-connect.h", + "lib/cf-ip-happy.c", + "lib/cf-ip-happy.h", + "lib/cf-recvbuf.c", + "lib/cf-recvbuf.h", + "lib/cf-setup.c", + "lib/cf-setup.h", "lib/cf-socket.c", "lib/cf-socket.h", "lib/cfilters.c", "lib/cfilters.h", - "lib/config-amigaos.h", - "lib/config-dos.h", "lib/config-mac.h", "lib/config-os400.h", - "lib/config-plan9.h", "lib/config-riscos.h", "lib/config-win32.h", - "lib/config-win32ce.h", "lib/conncache.c", "lib/conncache.h", "lib/connect.c", @@ -87,16 +89,19 @@ cc_library( "lib/content_encoding.h", "lib/cookie.c", "lib/cookie.h", + "lib/creds.c", + "lib/creds.h", + "lib/cshutdn.c", + "lib/cshutdn.h", "lib/curl_addrinfo.c", "lib/curl_addrinfo.h", - "lib/curl_base64.h", "lib/curl_ctype.h", - "lib/curl_des.c", - "lib/curl_des.h", "lib/curl_endian.c", "lib/curl_endian.h", "lib/curl_fnmatch.c", "lib/curl_fnmatch.h", + "lib/curl_fopen.c", + "lib/curl_fopen.h", "lib/curl_get_line.c", "lib/curl_get_line.h", "lib/curl_gethostname.c", @@ -104,45 +109,78 @@ cc_library( "lib/curl_gssapi.c", "lib/curl_gssapi.h", "lib/curl_hmac.h", - "lib/curl_krb5.h", "lib/curl_ldap.h", "lib/curl_md4.h", "lib/curl_md5.h", - "lib/curl_memory.h", "lib/curl_memrchr.c", "lib/curl_memrchr.h", - "lib/curl_multibyte.c", - "lib/curl_multibyte.h", "lib/curl_ntlm_core.c", "lib/curl_ntlm_core.h", "lib/curl_printf.h", "lib/curl_range.c", "lib/curl_range.h", - "lib/curl_rtmp.c", - "lib/curl_rtmp.h", "lib/curl_sasl.c", "lib/curl_sasl.h", "lib/curl_setup.h", - "lib/curl_setup_once.h", "lib/curl_sha256.h", "lib/curl_sha512_256.c", "lib/curl_sha512_256.h", + "lib/curl_share.c", + "lib/curl_share.h", "lib/curl_sspi.c", "lib/curl_sspi.h", "lib/curl_threads.c", "lib/curl_threads.h", "lib/curl_trc.c", "lib/curl_trc.h", - "lib/curlx.h", + "lib/curlx/base64.c", + "lib/curlx/base64.h", + "lib/curlx/basename.c", + "lib/curlx/basename.h", + "lib/curlx/dynbuf.c", + "lib/curlx/dynbuf.h", + "lib/curlx/fopen.c", + "lib/curlx/fopen.h", + "lib/curlx/inet_ntop.c", + "lib/curlx/inet_ntop.h", + "lib/curlx/inet_pton.c", + "lib/curlx/inet_pton.h", + "lib/curlx/multibyte.c", + "lib/curlx/multibyte.h", + "lib/curlx/nonblock.c", + "lib/curlx/nonblock.h", + "lib/curlx/snprintf.c", + "lib/curlx/snprintf.h", + "lib/curlx/strcopy.c", + "lib/curlx/strcopy.h", + "lib/curlx/strdup.c", + "lib/curlx/strdup.h", + "lib/curlx/strerr.c", + "lib/curlx/strerr.h", + "lib/curlx/strparse.c", + "lib/curlx/strparse.h", + "lib/curlx/timediff.c", + "lib/curlx/timediff.h", + "lib/curlx/timeval.c", + "lib/curlx/timeval.h", + "lib/curlx/version_win32.c", + "lib/curlx/version_win32.h", + "lib/curlx/wait.c", + "lib/curlx/wait.h", + "lib/curlx/warnless.c", + "lib/curlx/warnless.h", + "lib/curlx/winapi.c", + "lib/curlx/winapi.h", "lib/cw-out.c", "lib/cw-out.h", + "lib/cw-pause.c", + "lib/cw-pause.h", "lib/dict.c", "lib/dict.h", - "lib/dllmain.c", + "lib/dnscache.c", + "lib/dnscache.h", "lib/doh.c", "lib/doh.h", - "lib/dynbuf.c", - "lib/dynbuf.h", "lib/dynhds.c", "lib/dynhds.h", "lib/easy.c", @@ -153,16 +191,17 @@ cc_library( "lib/easyoptions.h", "lib/escape.c", "lib/escape.h", + "lib/fake_addrinfo.c", + "lib/fake_addrinfo.h", "lib/file.c", "lib/file.h", "lib/fileinfo.c", "lib/fileinfo.h", - "lib/fopen.c", - "lib/fopen.h", "lib/formdata.c", "lib/formdata.h", "lib/ftp.c", "lib/ftp.h", + "lib/ftp-int.h", "lib/ftplistparser.c", "lib/ftplistparser.h", "lib/functypes.h", @@ -176,12 +215,10 @@ cc_library( "lib/headers.c", "lib/headers.h", "lib/hmac.c", - "lib/hostasyn.c", "lib/hostip.c", "lib/hostip.h", "lib/hostip4.c", "lib/hostip6.c", - "lib/hostsyn.c", "lib/hsts.c", "lib/hsts.h", "lib/http.c", @@ -202,16 +239,14 @@ cc_library( "lib/http_ntlm.h", "lib/http_proxy.c", "lib/http_proxy.h", + "lib/httpsrr.c", + "lib/httpsrr.h", "lib/idn.c", "lib/idn.h", "lib/if2ip.c", "lib/if2ip.h", "lib/imap.c", "lib/imap.h", - "lib/inet_ntop.h", - "lib/inet_pton.c", - "lib/inet_pton.h", - "lib/krb5.c", "lib/ldap.c", "lib/llist.c", "lib/llist.h", @@ -220,36 +255,41 @@ cc_library( "lib/md4.c", "lib/md5.c", "lib/memdebug.c", - "lib/memdebug.h", "lib/mime.c", "lib/mime.h", "lib/mprintf.c", "lib/mqtt.c", "lib/mqtt.h", "lib/multi.c", + "lib/multi_ev.c", + "lib/multi_ev.h", + "lib/multi_ntfy.c", + "lib/multi_ntfy.h", "lib/multihandle.h", "lib/multiif.h", "lib/netrc.c", "lib/netrc.h", - "lib/nonblock.c", - "lib/nonblock.h", - "lib/noproxy.c", - "lib/noproxy.h", "lib/openldap.c", "lib/parsedate.c", "lib/parsedate.h", + "lib/peer.c", + "lib/peer.h", "lib/pingpong.c", "lib/pingpong.h", "lib/pop3.c", "lib/pop3.h", "lib/progress.c", "lib/progress.h", + "lib/protocol.c", + "lib/protocol.h", + "lib/proxy.c", + "lib/proxy.h", "lib/psl.c", "lib/psl.h", "lib/rand.c", "lib/rand.h", - "lib/rename.c", - "lib/rename.h", + "lib/ratelimit.c", + "lib/ratelimit.h", "lib/request.c", "lib/request.h", "lib/rtsp.c", @@ -262,9 +302,8 @@ cc_library( "lib/setopt.h", "lib/setup-os400.h", "lib/setup-vms.h", + "lib/setup-win32.h", "lib/sha256.c", - "lib/share.c", - "lib/share.h", "lib/sigpipe.h", "lib/slist.c", "lib/slist.h", @@ -279,31 +318,32 @@ cc_library( "lib/socks.h", "lib/socks_gssapi.c", "lib/socks_sspi.c", - "lib/speedcheck.c", - "lib/speedcheck.h", "lib/splay.c", "lib/splay.h", "lib/strcase.c", "lib/strcase.h", - "lib/strdup.c", - "lib/strdup.h", + "lib/strequal.c", "lib/strerror.c", "lib/strerror.h", - "lib/strtok.c", - "lib/strtok.h", - "lib/strtoofft.c", - "lib/strtoofft.h", "lib/system_win32.h", "lib/telnet.c", "lib/telnet.h", "lib/tftp.c", "lib/tftp.h", - "lib/timediff.c", - "lib/timediff.h", - "lib/timeval.c", - "lib/timeval.h", + "lib/thrdpool.c", + "lib/thrdpool.h", + "lib/thrdqueue.c", + "lib/thrdqueue.h", "lib/transfer.c", "lib/transfer.h", + "lib/uint-bset.c", + "lib/uint-bset.h", + "lib/uint-hash.c", + "lib/uint-hash.h", + "lib/uint-spbset.c", + "lib/uint-spbset.h", + "lib/uint-table.c", + "lib/uint-table.h", "lib/url.c", "lib/url.h", "lib/urlapi.c", @@ -318,7 +358,6 @@ cc_library( "lib/vauth/krb5_gssapi.c", "lib/vauth/krb5_sspi.c", "lib/vauth/ntlm.c", - "lib/vauth/ntlm.h", "lib/vauth/ntlm_sspi.c", "lib/vauth/oauth2.c", "lib/vauth/spnego_gssapi.c", @@ -326,26 +365,30 @@ cc_library( "lib/vauth/vauth.c", "lib/vauth/vauth.h", "lib/version.c", - "lib/version_win32.c", - "lib/version_win32.h", - "lib/vquic/curl_msh3.c", - "lib/vquic/curl_msh3.h", - "lib/vquic/curl_ngtcp2.c", - "lib/vquic/curl_ngtcp2.h", - "lib/vquic/curl_osslq.h", - "lib/vquic/curl_quiche.c", - "lib/vquic/curl_quiche.h", + "lib/vquic/capsule.c", + "lib/vquic/capsule.h", + "lib/vquic/cf-capsule.c", + "lib/vquic/cf-capsule.h", + "lib/vquic/cf-ngtcp2.c", + "lib/vquic/cf-ngtcp2.h", + "lib/vquic/cf-ngtcp2-cmn.c", + "lib/vquic/cf-ngtcp2-cmn.h", + "lib/vquic/cf-ngtcp2-proxy.c", + "lib/vquic/cf-ngtcp2-proxy.h", + "lib/vquic/cf-quiche.c", + "lib/vquic/cf-quiche.h", "lib/vquic/vquic.c", "lib/vquic/vquic.h", + "lib/vquic/vquic-tls.c", + "lib/vquic/vquic-tls.h", "lib/vquic/vquic_int.h", - "lib/vssh/curl_path.c", - "lib/vssh/curl_path.h", "lib/vssh/libssh.c", "lib/vssh/libssh2.c", "lib/vssh/ssh.h", - "lib/vssh/wolfssh.c", - "lib/vtls/bearssl.c", - "lib/vtls/bearssl.h", + "lib/vssh/vssh.c", + "lib/vssh/vssh.h", + "lib/vtls/apple.c", + "lib/vtls/apple.h", "lib/vtls/cipher_suite.c", "lib/vtls/cipher_suite.h", "lib/vtls/gtls.c", @@ -356,8 +399,6 @@ cc_library( "lib/vtls/keylog.h", "lib/vtls/mbedtls.c", "lib/vtls/mbedtls.h", - "lib/vtls/mbedtls_threadlock.c", - "lib/vtls/mbedtls_threadlock.h", "lib/vtls/openssl.c", "lib/vtls/openssl.h", "lib/vtls/rustls.c", @@ -366,28 +407,24 @@ cc_library( "lib/vtls/schannel.h", "lib/vtls/schannel_int.h", "lib/vtls/schannel_verify.c", - "lib/vtls/sectransp.h", "lib/vtls/vtls.c", "lib/vtls/vtls.h", + "lib/vtls/vtls_config.c", + "lib/vtls/vtls_config.h", "lib/vtls/vtls_int.h", + "lib/vtls/vtls_scache.c", + "lib/vtls/vtls_scache.h", + "lib/vtls/vtls_spack.c", + "lib/vtls/vtls_spack.h", "lib/vtls/wolfssl.c", "lib/vtls/wolfssl.h", "lib/vtls/x509asn1.c", "lib/vtls/x509asn1.h", - "lib/warnless.c", - "lib/warnless.h", "lib/ws.c", "lib/ws.h", ] + select({ - "@xla//xla/tsl:macos": [ - "lib/vtls/sectransp.c", - ], - "@xla//xla/tsl:ios": [ - "lib/vtls/sectransp.c", - ], "@xla//xla/tsl:windows": CURL_WIN_SRCS, - "//conditions:default": [ - ], + "//conditions:default": [], }), hdrs = [ "include/curl/curl.h", @@ -446,7 +483,9 @@ cc_library( "@xla//xla/tsl:windows": [ "-DEFAULTLIB:ws2_32.lib", "-DEFAULTLIB:advapi32.lib", + "-DEFAULTLIB:bcrypt.lib", "-DEFAULTLIB:crypt32.lib", + "-DEFAULTLIB:secur32.lib", "-DEFAULTLIB:Normaliz.lib", ], "//conditions:default": [ @@ -476,14 +515,14 @@ cc_binary( name = "curl_bin", srcs = [ "lib/config-win32.h", + ] + [ + "src/config2setopts.c", + "src/config2setopts.h", + "src/curlinfo.c", "src/slist_wc.c", "src/slist_wc.h", "src/terminal.c", "src/terminal.h", - "src/tool_binmode.c", - "src/tool_binmode.h", - "src/tool_bname.c", - "src/tool_bname.h", "src/tool_cb_dbg.c", "src/tool_cb_dbg.h", "src/tool_cb_hdr.c", @@ -508,6 +547,8 @@ cc_binary( "src/tool_easysrc.h", "src/tool_filetime.c", "src/tool_filetime.h", + "src/tool_findfile.c", + "src/tool_findfile.h", "src/tool_formparse.c", "src/tool_formparse.h", "src/tool_getparam.c", @@ -518,26 +559,21 @@ cc_binary( "src/tool_help.h", "src/tool_helpers.c", "src/tool_helpers.h", - "src/tool_homedir.c", - "src/tool_homedir.h", "src/tool_hugehelp.c", "src/tool_hugehelp.h", + "src/tool_ipfs.c", + "src/tool_ipfs.h", "src/tool_libinfo.c", "src/tool_libinfo.h", + "src/tool_listhelp.c", "src/tool_main.c", "src/tool_main.h", - "src/tool_metalink.c", - "src/tool_metalink.h", - "src/tool_mfiles.c", - "src/tool_mfiles.h", "src/tool_msgs.c", "src/tool_msgs.h", "src/tool_operate.c", "src/tool_operate.h", "src/tool_operhlp.c", "src/tool_operhlp.h", - "src/tool_panykey.c", - "src/tool_panykey.h", "src/tool_paramhlp.c", "src/tool_paramhlp.h", "src/tool_parsecfg.c", @@ -548,10 +584,10 @@ cc_binary( "src/tool_setopt.c", "src/tool_setopt.h", "src/tool_setup.h", - "src/tool_sleep.c", - "src/tool_sleep.h", - "src/tool_strdup.c", - "src/tool_strdup.h", + "src/tool_ssls.c", + "src/tool_ssls.h", + "src/tool_stderr.c", + "src/tool_stderr.h", "src/tool_urlglob.c", "src/tool_urlglob.h", "src/tool_util.c", @@ -559,14 +595,14 @@ cc_binary( "src/tool_version.h", "src/tool_vms.c", "src/tool_vms.h", - "src/tool_writeenv.c", - "src/tool_writeenv.h", "src/tool_writeout.c", "src/tool_writeout.h", "src/tool_writeout_json.c", "src/tool_writeout_json.h", "src/tool_xattr.c", "src/tool_xattr.h", + "src/var.c", + "src/var.h", ], copts = select({ "@xla//xla/tsl:windows": CURL_BIN_WIN_COPTS, @@ -622,7 +658,7 @@ genrule( "# define HAVE_SYS_FILIO_H 1", "# define HAVE_SYS_SOCKIO_H 1", "# define CURL_OS \"x86_64-apple-darwin15.5.0\"", - "# define USE_SECTRANSP 1", + "# define USE_APPLE_SECTRUST 1", "#else", "# define CURL_CA_BUNDLE \"/etc/ssl/certs/ca-certificates.crt\"", "# define GETSERVBYPORT_R_ARGS 6", diff --git a/third_party/xla/workspace2.bzl b/third_party/xla/workspace2.bzl index 0f49a1db57c284..d3d5aa85cec30d 100644 --- a/third_party/xla/workspace2.bzl +++ b/third_party/xla/workspace2.bzl @@ -371,9 +371,9 @@ def _tf_repositories(): tf_http_archive( name = "curl", build_file = "//third_party:curl.BUILD", - sha256 = "264537d90e58d2b09dddc50944baf3c38e7089151c8986715e2aaeaaf2b8118f", - strip_prefix = "curl-8.11.0", - urls = tf_mirror_urls("https://curl.se/download/curl-8.11.0.tar.gz"), + sha256 = "d9b327997999045a24cda50f3983e69e51c516bd8be6ef9842fc7f99135e33bb", + strip_prefix = "curl-8.21.0", + urls = tf_mirror_urls("https://curl.se/download/curl-8.21.0.tar.gz"), ) tf_http_archive( diff --git a/third_party/xla/xla/tsl/platform/cloud/curl_http_request.cc b/third_party/xla/xla/tsl/platform/cloud/curl_http_request.cc index c1fe888fa899e0..190fefb679194b 100644 --- a/third_party/xla/xla/tsl/platform/cloud/curl_http_request.cc +++ b/third_party/xla/xla/tsl/platform/cloud/curl_http_request.cc @@ -62,7 +62,9 @@ class LibCurlProxy : public LibCurl { CURLcode curl_easy_setopt(CURL* curl, CURLoption option, uint64_t param) override { - return ::curl_easy_setopt(curl, option, param); + // libcurl integer options take `long` (32-bit on Windows). Newer libcurl + // type-checks the C API, so convert at this boundary. + return ::curl_easy_setopt(curl, option, static_cast(param)); } CURLcode curl_easy_setopt(CURL* curl, CURLoption option, @@ -100,7 +102,12 @@ class LibCurlProxy : public LibCurl { CURLcode curl_easy_getinfo(CURL* curl, CURLINFO info, uint64_t* value) override { - return ::curl_easy_getinfo(curl, info, value); + long v = 0; + const CURLcode result = ::curl_easy_getinfo(curl, info, &v); + if (result == CURLE_OK) { + *value = static_cast(v); + } + return result; } CURLcode curl_easy_getinfo(CURL* curl, CURLINFO info, @@ -267,7 +274,7 @@ absl::Status CurlHttpRequest::SetPutFromFile(const std::string& body_filepath, curl_headers_ = libcurl_->curl_slist_append( curl_headers_, absl::StrCat("Content-Length: ", size).c_str()); - CHECK_CURL_OK(libcurl_->curl_easy_setopt(curl_, CURLOPT_PUT, 1)); + CHECK_CURL_OK(libcurl_->curl_easy_setopt(curl_, CURLOPT_UPLOAD, uint64_t{1})); CHECK_CURL_OK(libcurl_->curl_easy_setopt(curl_, CURLOPT_READDATA, reinterpret_cast(put_body_))); // Using the default CURLOPT_READFUNCTION, which is doing an fread() on the @@ -280,7 +287,7 @@ void CurlHttpRequest::SetPutEmptyBody() { CheckMethodNotSet(); is_method_set_ = true; method_ = RequestMethod::kPut; - CHECK_CURL_OK(libcurl_->curl_easy_setopt(curl_, CURLOPT_PUT, 1)); + CHECK_CURL_OK(libcurl_->curl_easy_setopt(curl_, CURLOPT_UPLOAD, uint64_t{1})); AddHeader("Content-Length", "0"); AddHeader("Transfer-Encoding", "identity"); CHECK_CURL_OK(libcurl_->curl_easy_setopt(curl_, CURLOPT_READDATA, @@ -462,10 +469,6 @@ absl::Status CurlHttpRequest::Send() { const CURLcode curl_result = libcurl_->curl_easy_perform(curl_); ABSL_RETURN_IF_ERROR(CURLcodeToStatus(curl_result, error_buffer)); - double written_size = 0; - CHECK_CURL_OK(libcurl_->curl_easy_getinfo(curl_, CURLINFO_SIZE_DOWNLOAD, - &written_size)); - CHECK_CURL_OK(libcurl_->curl_easy_getinfo(curl_, CURLINFO_RESPONSE_CODE, &response_code_)); diff --git a/third_party/xla/xla/tsl/platform/cloud/curl_http_request_test.cc b/third_party/xla/xla/tsl/platform/cloud/curl_http_request_test.cc index 77c1d2102b2956..d0ea2259a59f6c 100644 --- a/third_party/xla/xla/tsl/platform/cloud/curl_http_request_test.cc +++ b/third_party/xla/xla/tsl/platform/cloud/curl_http_request_test.cc @@ -68,7 +68,7 @@ class FakeLibCurl : public LibCurl { case CURLOPT_POST: is_post_ = param; break; - case CURLOPT_PUT: + case CURLOPT_UPLOAD: is_put_ = param; break; default: @@ -192,15 +192,8 @@ class FakeLibCurl : public LibCurl { } return CURLE_OK; } - CURLcode curl_easy_getinfo(CURL* curl, CURLINFO info, - double* value) override { - switch (info) { - case CURLINFO_SIZE_DOWNLOAD: - *value = response_content_.size(); - break; - default: - break; - } + CURLcode curl_easy_getinfo(CURL* /*curl*/, CURLINFO /*info*/, + double* /*value*/) override { return CURLE_OK; } void curl_easy_cleanup(CURL* curl) override { is_cleaned_up_ = true; } From 4be4b7e4febb01b2efb2c0b373116e5134508f3e Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Mon, 21 Sep 2026 07:11:19 -0700 Subject: [PATCH 05/17] Remove obsolete third_party/llvm/build.patch This patch disabled BLAKE3 x86_64 SIMD assembly in LLVM's Bazel overlay to work around a 2022 Windows build failure where MSVC's ml64.exe failed on Unix .S assembly files. Upstream LLVM now selects the Windows GNU assembly files for clang-cl, and XLA Windows builds use clang-cl instead of MSVC. Removing this patch also restores SIMD-accelerated BLAKE3 hashing on x86_64 platforms. PiperOrigin-RevId: 985256166 --- third_party/xla/third_party/llvm/build.patch | 39 ------------------- .../xla/third_party/llvm/workspace.bzl | 1 - 2 files changed, 40 deletions(-) delete mode 100644 third_party/xla/third_party/llvm/build.patch diff --git a/third_party/xla/third_party/llvm/build.patch b/third_party/xla/third_party/llvm/build.patch deleted file mode 100644 index 032406b4e16483..00000000000000 --- a/third_party/xla/third_party/llvm/build.patch +++ /dev/null @@ -1,39 +0,0 @@ -# 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. -# ============================================================================== -diff --git a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel b/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel ---- a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel -+++ b/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel -@@ -430,7 +430,8 @@ - "@platforms//cpu:aarch64": [ - "lib/Support/BLAKE3/blake3_neon.c", - ], -- "@platforms//cpu:x86_64": [":blake3_x86_64_asm"], -+ # TODO(b/234415414): temporary disabled -+ # "@platforms//cpu:x86_64": [":blake3_x86_64_asm"], - "//conditions:default": [ - ], - }), -@@ -459,8 +460,9 @@ - ], - "//conditions:default": ["BLAKE3_USE_NEON=0"], - }) + select({ -- "@platforms//cpu:x86_64": [ -- ], -+ # TODO(b/234415414): temporary disabled -+ # "@platforms//cpu:x86_64": [ -+ # ], - "//conditions:default": [ - "BLAKE3_NO_AVX2", - "BLAKE3_NO_AVX512", diff --git a/third_party/xla/third_party/llvm/workspace.bzl b/third_party/xla/third_party/llvm/workspace.bzl index 83d4a833e45971..08380e9e8ff06e 100644 --- a/third_party/xla/third_party/llvm/workspace.bzl +++ b/third_party/xla/third_party/llvm/workspace.bzl @@ -33,7 +33,6 @@ def repo(name): build_file = "//third_party/llvm:llvm.BUILD", patch_file = [ "//third_party/llvm:generated.patch", # Autogenerated, don't remove. - "//third_party/llvm:build.patch", "//third_party/llvm:toolchains.patch", "//third_party/llvm:lit_test.patch", "//third_party/llvm:run_lit.patch", From cbec5fc4f1118b18ec6333e2123ae67ef5dc8bce Mon Sep 17 00:00:00 2001 From: Shyamli Agrawal Date: Mon, 21 Sep 2026 07:55:54 -0700 Subject: [PATCH 06/17] Allow ConfigAssigner to handle cuBLASLt GEMM instructions even when xla_gpu_autotune_level is 0. - This particular condition of changing the instruction filter depending on the autotune level can lead to inconsistency and confusion. PiperOrigin-RevId: 985277541 --- .../xla/backends/gpu/autotuner/hipblaslt.cc | 16 ++++++- .../backends/gpu/autotuner/hipblaslt_test.cc | 5 ++- .../gpu/codegen/triton/dot_algorithms_test.cc | 13 +++--- .../gpu/autotuning/config_assigner_pass.cc | 6 +-- .../autotuning/config_assigner_pass_test.cc | 45 +++++++++++++++++++ 5 files changed, 74 insertions(+), 11 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/autotuner/hipblaslt.cc b/third_party/xla/xla/backends/gpu/autotuner/hipblaslt.cc index 9a0ec963f79d8e..4c3a6f071ab220 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/hipblaslt.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/hipblaslt.cc @@ -389,7 +389,21 @@ HipblasLtBackend::GetDefaultConfig(const HloInstruction& instr) { } auto config = std::make_unique(); - config->mutable_gemm()->set_algorithm(0); + auto* gemm_key = config->mutable_gemm(); + gemm_key->set_algorithm(0); + // Standard GEMMs were already lowered into custom calls with a workspace + // buffer tuple by GemmRewriter, so we can preserve that allocated workspace + // size. In contrast, scaled-dot GEMMs enter this pass as a fusion and are + // only lowered to a custom call in ApplyConfig(), so they do not yet have a + // workspace buffer in their shape. MX scaled dots on ROCm are currently + // supported on gfx950, so we use kGFX950Workspace (matching + // GetSupportedConfigs). + if (GetScaledDotFromFusion(instr) != nullptr) { + gemm_key->set_autotune_workspace_size(GemmConfig::kGFX950Workspace); + } else if (instr.shape().IsTuple() && !instr.shape().tuple_shapes().empty()) { + gemm_key->set_autotune_workspace_size( + ShapeUtil::ByteSizeOf(instr.shape().tuple_shapes().back())); + } return config; } diff --git a/third_party/xla/xla/backends/gpu/autotuner/hipblaslt_test.cc b/third_party/xla/xla/backends/gpu/autotuner/hipblaslt_test.cc index d7ddd2f7aa2a19..ccad178efb4d90 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/hipblaslt_test.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/hipblaslt_test.cc @@ -172,7 +172,10 @@ TEST_F(HipblasLtBackendTest, GetDefaultConfig) { absl::StatusOr> config = backend_.GetDefaultConfig( (*module->entry_computation()->root_instruction()->operand(0))); - EXPECT_THAT(config, absl_testing::IsOk()); + ASSERT_THAT(config, absl_testing::IsOk()); + ASSERT_TRUE((*config)->has_gemm()); + EXPECT_THAT((*config)->gemm().algorithm(), 0); + EXPECT_NE((*config)->gemm().autotune_workspace_size(), 0); } TEST_F(HipblasLtBackendTest, GetDefaultConfigFailsWithoutAHipblasLtCustomCall) { diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc index 81a06dc5487b9f..21d02db5101fc1 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc @@ -465,18 +465,19 @@ TEST_F(BlasAlgorithmTest, Algorithm_TF32_TF32_F32_X3) { break; case CudaComputeCapabilities::kHopper: { DebugOptions debug_options = GetDebugOptionsForTest(); - std::string dot_kernel_name = "tf32f32"; + ::testing::Matcher dot_kernel_matcher = + ::testing::HasSubstr("tf32f32"); if (debug_options.xla_gpu_enable_cublaslt()) { - // CublasLt uses cutlass for TF32. - dot_kernel_name = "cutlass_80"; + dot_kernel_matcher = + ::testing::AnyOf(::testing::HasSubstr("cutlass_80"), + ::testing::HasSubstr("sm90_xmma_gemm")); } EXPECT_THAT(kernel_names, ::testing::UnorderedElementsAre( ::testing::HasSubstr("loop_and_subtract"), ::testing::HasSubstr("loop_and_subtract"), ::testing::HasSubstr("loop_select_fusion"), - ::testing::HasSubstr(dot_kernel_name), - ::testing::HasSubstr(dot_kernel_name), - ::testing::HasSubstr(dot_kernel_name))); + dot_kernel_matcher, dot_kernel_matcher, + dot_kernel_matcher)); break; } default: diff --git a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc index 48c994bb414166..67c4d9f647cade 100644 --- a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc +++ b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc @@ -361,8 +361,7 @@ InstructionFilterFn GetShouldAssignConfigToInstructionFn( const DebugOptions& debug_options, const se::GpuComputeCapability& gpu_version) { bool do_not_autotune_cublas = - debug_options.xla_gpu_experimental_disable_binary_libraries() || - debug_options.xla_gpu_autotune_level() == 0; + debug_options.xla_gpu_experimental_disable_binary_libraries(); bool do_not_autotune_cudnn = debug_options.xla_gpu_experimental_disable_binary_libraries() || (do_not_autotune_cublas && !gpu_version.IsRocm()); @@ -374,7 +373,8 @@ InstructionFilterFn GetShouldAssignConfigToInstructionFn( debug_options.xla_gpu_experimental_enable_fusion_autotuner(); return [do_not_autotune_cublas, do_not_autotune_cudnn, - enable_fusion_autotuner](const HloInstruction& instruction) -> bool { + enable_fusion_autotuner, + gpu_version](const HloInstruction& instruction) -> bool { AutotuneDecision decision = ShouldAssignConfigToInstruction( do_not_autotune_cublas, do_not_autotune_cudnn, enable_fusion_autotuner, instruction); diff --git a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc index fb99d9596a7ac7..f7ec78ec23a464 100644 --- a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc +++ b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc @@ -498,6 +498,51 @@ TEST_F(ConfigAssignerPassTest, DevicelessUsesDefaultConfigIfNoCache) { gpu_backend_config.gemm_backend_config().has_selected_algorithm()); } +TEST_F(ConfigAssignerPassTest, AutotuneLevel0UsesDefaultConfig) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kCublasCustomCallHlo)); + + module->mutable_config().mutable_debug_options().set_xla_gpu_autotune_level( + 0); + + tsl::thread::ThreadPool thread_pool(tsl::Env::Default(), "autotuning", + /*num_threads=*/4); + GpuCompiler::GpuTargetConfig target_config(stream_executor_); + + std::vector> backends; + backends.push_back(std::make_unique( + stream_executor_, &module->config().debug_options(), &compiler_, + &target_config)); + + auto get_backends_fn = + [backends = + std::make_shared>>( + std::move(backends))]() mutable { return std::move(*backends); }; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr pass, + ConfigAssignerPass::Create( + std::move(get_backends_fn), module->config().debug_options(), + target_config.device_description.gpu_compute_capability(), + stream_executor_, &thread_pool, &target_config, + /*alias_info=*/nullptr, /*mlir_context=*/nullptr, + /*shape_size_fn=*/[](const Shape& shape) { return 0; }, + allocator_.get())); + EXPECT_THAT(pass->Run(module.get(), /*execution_threads=*/{}), + absl_testing::IsOkAndHolds(true)); + + // Verify that the backend config has been updated in the HLO with default + // config. + auto gemm = + module->entry_computation()->GetInstructionWithName("custom-call.1"); + ASSERT_OK_AND_ASSIGN(auto gpu_backend_config, + gemm->backend_config()); + EXPECT_TRUE( + gpu_backend_config.gemm_backend_config().has_selected_algorithm()); + EXPECT_EQ(gpu_backend_config.gemm_backend_config().selected_algorithm(), 0); + EXPECT_EQ(gpu_backend_config.gemm_backend_config().autotune_workspace_size(), + 80000); +} + TEST_F(ConfigAssignerPassTest, CublasGemmInNonDefaultStreamIsAutotuned) { const char kCublasCustomNonDefaultStreamCallHlo[] = R"""( HloModule module, entry_computation_layout={(f32[100,100]{1,0}, f32[100,100]{1,0})->f32[100,100]{1,0}} From 4a78fda60a1384aab28a27089a8174bca2748fa8 Mon Sep 17 00:00:00 2001 From: Peter Hawkins Date: Mon, 21 Sep 2026 08:19:09 -0700 Subject: [PATCH 07/17] [XLA] Upgrade nanobind to v3.1.0 PiperOrigin-RevId: 985290320 --- third_party/xla/third_party/nanobind/workspace.bzl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/third_party/xla/third_party/nanobind/workspace.bzl b/third_party/xla/third_party/nanobind/workspace.bzl index 3370d0f71246c1..5a265d0b274a43 100644 --- a/third_party/xla/third_party/nanobind/workspace.bzl +++ b/third_party/xla/third_party/nanobind/workspace.bzl @@ -20,8 +20,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls") def repo(): tf_http_archive( name = "nanobind", - strip_prefix = "nanobind-db4827f06f6f1680e5d4004c95fc8d69299dba8b", - sha256 = "60c350b2d64cdf2c1d8e46433d38dba0fb69ba118f9077c1c7e74228e6dec1e0", - urls = tf_mirror_urls("https://github.com/wjakob/nanobind/archive/db4827f06f6f1680e5d4004c95fc8d69299dba8b.tar.gz"), + strip_prefix = "nanobind-68480a9e6883bb1bf65e1925423321e2a89d07f9", + sha256 = "177ca257aacbf58004809f8cd1c20ad0416d8ace0a884c16a3c03ea87e81876b", + urls = tf_mirror_urls("https://github.com/wjakob/nanobind/archive/68480a9e6883bb1bf65e1925423321e2a89d07f9.tar.gz"), build_file = "//third_party/nanobind:nanobind.BUILD", ) From 859147b1237072daa7260415eec5bc4bf8dca683 Mon Sep 17 00:00:00 2001 From: Joshua Lang Date: Mon, 21 Sep 2026 09:32:44 -0700 Subject: [PATCH 08/17] Move CPU-only dynamic_update_slice/out_of_place.hlo test to dedicated cpu_only_tests lit_test_suite PiperOrigin-RevId: 985329328 --- .../xla/xla/codegen/emitters/tests/BUILD | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/third_party/xla/xla/codegen/emitters/tests/BUILD b/third_party/xla/xla/codegen/emitters/tests/BUILD index 2a0f84a1a1fdae..df8298a20dc381 100644 --- a/third_party/xla/xla/codegen/emitters/tests/BUILD +++ b/third_party/xla/xla/codegen/emitters/tests/BUILD @@ -48,11 +48,15 @@ _FILECHECK_ONLY_TESTS = [ "concatenate/test_small_dim.hlo", ] +_CPU_ONLY_TESTS = [ + "dynamic_update_slice/out_of_place.hlo", +] + lit_test_suite( name = "tests", srcs = glob( ["**/*.hlo"], - exclude = _FILECHECK_ONLY_TESTS, + exclude = _FILECHECK_ONLY_TESTS + _CPU_ONLY_TESTS, ), cfg = "//xla:lit.cfg.py", default_tags = tf_cuda_tests_tags(), @@ -79,3 +83,15 @@ lit_test_suite( "@llvm-project//llvm:FileCheck", ], ) + +lit_test_suite( + name = "cpu_only_tests", + srcs = _CPU_ONLY_TESTS, + cfg = "//xla:lit.cfg.py", + tools = [ + "//xla/backends/cpu/codegen/tools:cpu_fusion_to_mlir", + "//xla/backends/cpu/codegen/tools:cpu_test_correctness", + "//xla/codegen/tools:emitters_opt", + "@llvm-project//llvm:FileCheck", + ], +) From 5c4aee3ee66e984ef8d7c0b62008abaa7743c3fe Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 21 Sep 2026 09:38:26 -0700 Subject: [PATCH 09/17] Add PjRtIsCompilerVariantRegistered to query the PjRt compiler registry PiperOrigin-RevId: 985332763 --- third_party/xla/xla/pjrt/pjrt_compiler.cc | 19 ++++++++++++ third_party/xla/xla/pjrt/pjrt_compiler.h | 13 +++++++++ .../xla/xla/pjrt/pjrt_compiler_test.cc | 29 +++++++++++++++++++ 3 files changed, 61 insertions(+) diff --git a/third_party/xla/xla/pjrt/pjrt_compiler.cc b/third_party/xla/xla/pjrt/pjrt_compiler.cc index d92c0252d4bc0a..ac790d4b5be657 100644 --- a/third_party/xla/xla/pjrt/pjrt_compiler.cc +++ b/third_party/xla/xla/pjrt/pjrt_compiler.cc @@ -125,6 +125,19 @@ absl::StatusOr PjRtCompilerRegistry::GetCompiler( return GetOrCreateCompiler(platform_name, variant_name); } +bool PjRtCompilerRegistry::IsCompilerRegistered( + absl::string_view platform_name, absl::string_view variant_name) { + PjRtCompilerType key{platform_name, variant_name}; + { + absl::MutexLock l(compiler_mutex_); + if (compilers_.contains(key)) { + return true; + } + } + absl::MutexLock l(factory_mutex_); + return factories_.contains(key); +} + absl::Status PjRtCompilerRegistry::InitializeVariant( absl::string_view platform_name, absl::string_view variant_name) { return GetOrCreateCompiler(platform_name, variant_name).status(); @@ -186,6 +199,12 @@ absl::Status PjRtInitializeCompilerVariants() { return PjRtCompilerRegistry::Global().InitializeAllVariants(); } +bool PjRtIsCompilerVariantRegistered(absl::string_view platform_name, + absl::string_view variant_name) { + return PjRtCompilerRegistry::Global().IsCompilerRegistered(platform_name, + variant_name); +} + void PjRtRegisterDefaultCompiler(absl::string_view platform_name, std::unique_ptr compiler) { CHECK_OK(PjRtCompilerRegistry::Global().RegisterCompiler( diff --git a/third_party/xla/xla/pjrt/pjrt_compiler.h b/third_party/xla/xla/pjrt/pjrt_compiler.h index 239af2e09ceacd..e450ff8eefeb4f 100644 --- a/third_party/xla/xla/pjrt/pjrt_compiler.h +++ b/third_party/xla/xla/pjrt/pjrt_compiler.h @@ -181,6 +181,13 @@ class PjRtCompilerRegistry { absl::StatusOr GetCompiler(absl::string_view platform_name, absl::string_view variant_name); + // Returns true if a compiler instance or a compiler factory is registered + // for the given platform and variant. Unlike GetCompiler(), this never + // instantiates the compiler. + bool IsCompilerRegistered(absl::string_view platform_name, + absl::string_view variant_name) + ABSL_LOCKS_EXCLUDED(compiler_mutex_, factory_mutex_); + // Explicitly initializes a compiler with a given variant. absl::Status InitializeVariant(absl::string_view platform_name, absl::string_view variant_name); @@ -246,6 +253,12 @@ absl::Status PjRtInitializeCompilerVariant(absl::string_view platform_name, // Initializes all compiler variants. absl::Status PjRtInitializeCompilerVariants(); +// Returns true if a compiler or a compiler factory is registered +// for the given platform and variant, i.e. if the variant can be served by +// this binary. Does not instantiate the compiler. +bool PjRtIsCompilerVariantRegistered(absl::string_view platform_name, + absl::string_view variant_name); + class PjRtClient; // Abstract interface to represent device topology that is used by the compiler. diff --git a/third_party/xla/xla/pjrt/pjrt_compiler_test.cc b/third_party/xla/xla/pjrt/pjrt_compiler_test.cc index f36c1f4e695619..ad19d44deeea8a 100644 --- a/third_party/xla/xla/pjrt/pjrt_compiler_test.cc +++ b/third_party/xla/xla/pjrt/pjrt_compiler_test.cc @@ -278,6 +278,35 @@ TEST(PjRtCompilerTest, VariantRegistryLookup) { EXPECT_TRUE(absl::IsNotFound(status.status())); } +TEST(PjRtCompilerTest, IsCompilerVariantRegistered) { + const std::string platform = "has_variant_test_platform"; + const std::string factory_variant = "factory_registered_variant"; + const std::string compiler_variant = "compiler_registered_variant"; + + EXPECT_FALSE(PjRtIsCompilerVariantRegistered(platform, factory_variant)); + EXPECT_FALSE(PjRtIsCompilerVariantRegistered(platform, compiler_variant)); + + bool factory_called = false; + PjRtRegisterCompilerFactory( + platform, factory_variant, + [&factory_called]() -> absl::StatusOr> { + factory_called = true; + return std::make_unique(); + }); + PjRtRegisterCompiler(platform, compiler_variant, + std::make_unique()); + + EXPECT_TRUE(PjRtIsCompilerVariantRegistered(platform, factory_variant)); + EXPECT_TRUE(PjRtIsCompilerVariantRegistered(platform, compiler_variant)); + // The query must not instantiate the compiler. + EXPECT_FALSE(factory_called); + + EXPECT_FALSE( + PjRtIsCompilerVariantRegistered(platform, "unregistered_variant")); + EXPECT_FALSE( + PjRtIsCompilerVariantRegistered("wrong_platform", factory_variant)); +} + TEST(PjRtTopologyDescriptionTest, DefaultMemorySpaceKindIds) { PjRtTestTopology topology; EXPECT_THAT(topology.GetMemorySpaceKindIds(), ::testing::ElementsAre(-1)); From ebab5d08c89b8b99f20f7ce37cbe40bb33daa381 Mon Sep 17 00:00:00 2001 From: David Pizzuto Date: Mon, 21 Sep 2026 09:46:41 -0700 Subject: [PATCH 10/17] mlir_to_hlo: Allow serializing to any Riegeli writer, not just strings. Riegeli's Writer abstraction doesn't imply any compression, so this simply allows for flexibility in the storage format, such as to Cord rather than string. In order to do so, we need to create a quick shim between Writer and llvm's raw_ostream, which handles the buffering and flushing already. PiperOrigin-RevId: 985337262 --- third_party/xla/xla/pjrt/BUILD | 3 + third_party/xla/xla/pjrt/mlir_to_hlo.cc | 66 +++++++++++++++++--- third_party/xla/xla/pjrt/mlir_to_hlo.h | 16 +++++ third_party/xla/xla/pjrt/mlir_to_hlo_test.cc | 23 +++++++ 4 files changed, 101 insertions(+), 7 deletions(-) diff --git a/third_party/xla/xla/pjrt/BUILD b/third_party/xla/xla/pjrt/BUILD index f48c112bb5678b..e471f9f745a0cc 100644 --- a/third_party/xla/xla/pjrt/BUILD +++ b/third_party/xla/xla/pjrt/BUILD @@ -819,6 +819,8 @@ cc_library( "@llvm-project//mlir:ShapeDialect", "@llvm-project//mlir:Support", "@llvm-project//mlir:Transforms", + "@riegeli//riegeli/bytes:string_writer", + "@riegeli//riegeli/bytes:writer", "@shardy//shardy/dialect/sdy/ir:dialect", "@shardy//shardy/dialect/sdy/ir:register", "@stablehlo//:chlo_ops", @@ -844,6 +846,7 @@ xla_cc_test( "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest_main", "@llvm-project//mlir:IR", + "@riegeli//riegeli/bytes:string_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 2f6d96131c426a..2fc7f207516248 100644 --- a/third_party/xla/xla/pjrt/mlir_to_hlo.cc +++ b/third_party/xla/xla/pjrt/mlir_to_hlo.cc @@ -15,6 +15,8 @@ limitations under the License. #include "xla/pjrt/mlir_to_hlo.h" +#include +#include #include #include #include @@ -58,6 +60,8 @@ limitations under the License. #include "shardy/dialect/sdy/ir/compatibility.h" #include "shardy/dialect/sdy/ir/dialect.h" #include "shardy/dialect/sdy/ir/register.h" +#include "riegeli/bytes/string_writer.h" +#include "riegeli/bytes/writer.h" #include "stablehlo/api/PortableApi.h" #include "stablehlo/dialect/ChloOps.h" #include "stablehlo/dialect/Register.h" @@ -82,8 +86,30 @@ limitations under the License. namespace xla { namespace { + using mlir::mhlo::ChloLegalizeToHighLevelMhloPassOptions; -} + +class RiegeliRawOstream : public llvm::raw_ostream { + public: + explicit RiegeliRawOstream(riegeli::Writer* writer) : writer_(writer) { + SetUnbuffered(); + } + ~RiegeliRawOstream() override = default; + + private: + uint64_t current_pos() const override { return writer_->pos(); } + + void write_impl(const char* ptr, size_t size) override { + if (size == 0) { + return; + } + writer_->Write(absl::string_view(ptr, size)); + } + + riegeli::Writer* writer_; +}; + +} // namespace void RegisterAllHloDialects(mlir::DialectRegistry& registry) { registry.insert(); @@ -286,10 +312,15 @@ absl::StatusOr ExpectSuccess(mlir::FailureOr result, std::string msg) { return *result; } -absl::StatusOr SerializeUsingVersionedStablehlo( - mlir::ModuleOp mlir_module, absl::string_view requested_target, - absl::string_view sdy_version, bool inplace, - bool allow_mixed_serialization) { +absl::Status SerializeToRiegeli(mlir::ModuleOp mlir_module, + absl::string_view requested_target, + absl::string_view sdy_version, bool inplace, + bool allow_mixed_serialization, + riegeli::Writer* writer) { + if (writer == nullptr) { + return absl::InvalidArgumentError("Writer cannot be null."); + } + mlir::MLIRContext* context = mlir_module->getContext(); mlir::BaseScopedDiagnosticHandler diagnostic_handler(context); @@ -373,8 +404,7 @@ absl::StatusOr SerializeUsingVersionedStablehlo( } // Serialize portable artifact - std::string buffer; - llvm::raw_string_ostream os(buffer); + RiegeliRawOstream os(writer); if (mlir::failed(mlir::stablehlo::serializePortableArtifact( mlir_module, target, os, /*allowOtherDialects=*/allow_mixed_serialization))) { @@ -383,6 +413,21 @@ absl::StatusOr SerializeUsingVersionedStablehlo( absl::StrCat("Failed to serialize StableHLO to plugin version ", target, ";\n\nDetailed error from MLIR: ", status.message())); } + return absl::OkStatus(); +} + +absl::StatusOr SerializeUsingVersionedStablehlo( + mlir::ModuleOp mlir_module, absl::string_view requested_target, + absl::string_view sdy_version, bool inplace, + bool allow_mixed_serialization) { + std::string buffer; + riegeli::StringWriter<> writer(&buffer); + ABSL_RETURN_IF_ERROR(SerializeToRiegeli(mlir_module, requested_target, sdy_version, + inplace, allow_mixed_serialization, + &writer)); + if (!writer.Close()) { + return writer.status(); + } return buffer; } @@ -426,6 +471,13 @@ absl::StatusOr Serialize(mlir::ModuleOp module, /*allow_mixed_serialization=*/true); } +absl::Status Serialize(mlir::ModuleOp module, absl::string_view target, + absl::string_view sdy_version, riegeli::Writer* writer, + bool inplace) { + return SerializeToRiegeli(module, target, sdy_version, inplace, + /*allow_mixed_serialization=*/true, writer); +} + absl::StatusOr Serialize(mlir::ModuleOp mlir_module, absl::string_view target, bool inplace) { return Serialize(mlir_module, target, GetDefaultSdyVersion(), inplace); diff --git a/third_party/xla/xla/pjrt/mlir_to_hlo.h b/third_party/xla/xla/pjrt/mlir_to_hlo.h index e4002bec906064..8397c3974dba3d 100644 --- a/third_party/xla/xla/pjrt/mlir_to_hlo.h +++ b/third_party/xla/xla/pjrt/mlir_to_hlo.h @@ -27,6 +27,7 @@ limitations under the License. #include "mlir/IR/MLIRContext.h" #include "mlir/IR/OwningOpRef.h" #include "mlir/Support/LLVM.h" +#include "riegeli/bytes/writer.h" #include "xla/client/executable_build_options.h" #include "xla/hlo/builder/xla_computation.h" #include "xla/mlir_hlo/mhlo/transforms/passes.h" @@ -90,6 +91,12 @@ absl::StatusOr Serialize(mlir::ModuleOp mlir_module, absl::string_view sdy_version, bool inplace = false); +// Serializes an MLIR module to a portable artifact, directing bytes directly +// into a riegeli::Writer. +absl::Status Serialize(mlir::ModuleOp mlir_module, absl::string_view target, + absl::string_view sdy_version, riegeli::Writer* writer, + bool inplace = false); + // Temporary legacy version of Serialize. Will use the default SDY version. // TODO(hyeontaek): Delete once all users are migrated. absl::StatusOr Serialize(mlir::ModuleOp mlir_module, @@ -126,6 +133,15 @@ absl::StatusOr SerializeUsingVersionedStablehlo( absl::string_view sdy_version, bool inplace = false, bool allow_mixed_serialization = false); +// Serializes an MLIR module to a portable artifact directing bytes directly +// into the provided riegeli::Writer, avoiding std::string contiguous +// allocation. +absl::Status SerializeToRiegeli(mlir::ModuleOp mlir_module, + absl::string_view requested_target, + absl::string_view sdy_version, bool inplace, + bool allow_mixed_serialization, + riegeli::Writer* writer); + // Temporary legacy version of SerializeUsingVersionedStablehlo. Will use // the default SDY version. // TODO(hyeontaek): Delete once all users are migrated. 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 2066888308654f..63d66f04e5611b 100644 --- a/third_party/xla/xla/pjrt/mlir_to_hlo_test.cc +++ b/third_party/xla/xla/pjrt/mlir_to_hlo_test.cc @@ -26,6 +26,7 @@ limitations under the License. #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/MLIRContext.h" #include "mlir/IR/OwningOpRef.h" +#include "riegeli/bytes/string_writer.h" #include "stablehlo/api/PortableApi.h" #include "xla/hlo/testlib/test.h" #include "xla/tsl/platform/statusor.h" @@ -261,5 +262,27 @@ TEST(MlirToHloTest, InvalidBytecodeTest) { HasSubstr(mlir::stablehlo::getCurrentVersion())); } +TEST(MlirToHloTest, SerializeToRiegeliTest) { + 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)); + + std::string buffer; + riegeli::StringWriter<> writer(&buffer); + EXPECT_OK(SerializeToRiegeli(*module, /*requested_target=*/"1.0.0", + /*sdy_version=*/"0.0.1", /*inplace=*/false, + /*allow_mixed_serialization=*/true, &writer)); + EXPECT_TRUE(writer.Close()); + EXPECT_THAT(buffer, IsVhloArtifact("1.0.0")); +} + } // namespace } // namespace xla From b201fe717ee7ad375953a8e19fe79d233abee737 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 21 Sep 2026 09:58:58 -0700 Subject: [PATCH 11/17] Add xla_cpu_mock_custom_calls debug option to DebugOptions and CPU ThunkEmitter, allowing AOT/offline compilation of HLO modules with custom calls whose FFI handlers are not registered in the compilation binary (matching the behavior of xla_gpu_mock_custom_calls on GPU). - Add xla_cpu_mock_custom_calls field to DebugOptions in xla.proto and debug_options_flags.cc. - In CPU ThunkEmitter::EmitCustomCallThunk, ignore NOT_FOUND errors for custom call handler lookups when xla_cpu_mock_custom_calls is enabled. - Add mock_custom_call_test unit test for CPU. Testing: Added mock_custom_call_test unit test for CPU and tested compilation of HLO modules containing custom calls. PiperOrigin-RevId: 985344122 --- third_party/xla/xla/backends/cpu/tests/BUILD | 11 ++++ .../cpu/tests/mock_custom_call_test.cc | 60 +++++++++++++++++++ third_party/xla/xla/debug_options_flags.cc | 5 ++ .../xla/xla/service/cpu/thunk_emitter.cc | 18 +++++- third_party/xla/xla/xla.proto | 3 + 5 files changed, 94 insertions(+), 3 deletions(-) create mode 100644 third_party/xla/xla/backends/cpu/tests/mock_custom_call_test.cc diff --git a/third_party/xla/xla/backends/cpu/tests/BUILD b/third_party/xla/xla/backends/cpu/tests/BUILD index 79c4519867f9be..e4e5e3541f4f8a 100644 --- a/third_party/xla/xla/backends/cpu/tests/BUILD +++ b/third_party/xla/xla/backends/cpu/tests/BUILD @@ -65,3 +65,14 @@ xla_test( "@com_google_googletest//:gtest_main", ], ) + +xla_test( + name = "mock_custom_call_test", + srcs = ["mock_custom_call_test.cc"], + backends = ["cpu"], + deps = [ + "//xla:xla_proto_cc", + "//xla/tests:hlo_test_base", + "@com_google_googletest//:gtest_main", + ], +) diff --git a/third_party/xla/xla/backends/cpu/tests/mock_custom_call_test.cc b/third_party/xla/xla/backends/cpu/tests/mock_custom_call_test.cc new file mode 100644 index 00000000000000..cf1d9a99488b19 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/tests/mock_custom_call_test.cc @@ -0,0 +1,60 @@ +/* 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 "xla/tests/hlo_test_base.h" +#include "xla/xla.pb.h" + +namespace xla::cpu { +namespace { + +using UnknownCustomCallFails = HloTestBase; + +TEST_F(UnknownCustomCallFails, UnknownCustomCallFails) { + const char* hlo_text = R"( + HloModule Test1 + + ENTRY Test1 { + a = f32[128] parameter(0) + ROOT r1 = f32[128] custom-call(a), custom_call_target="my_custom_call", api_version=API_VERSION_TYPED_FFI + } + )"; + + EXPECT_FALSE(Run(hlo_text)); +} + +class MockedCustomCall : public HloTestBase { + DebugOptions GetDebugOptionsForTest() const override { + DebugOptions opts = HloTestBase::GetDebugOptionsForTest(); + opts.set_xla_cpu_mock_custom_calls(true); + return opts; + } +}; + +TEST_F(MockedCustomCall, CustomCallIgnored) { + const char* hlo_text = R"( + HloModule Test1 + + ENTRY Test1 { + a = f32[128] parameter(0) + ROOT r1 = f32[128] custom-call(a), custom_call_target="my_custom_call", api_version=API_VERSION_TYPED_FFI + } + )"; + + EXPECT_TRUE(Run(hlo_text)); +} + +} // namespace +} // namespace xla::cpu diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc index cb8a39fa3e23d5..96a3de8ba6f1df 100644 --- a/third_party/xla/xla/debug_options_flags.cc +++ b/third_party/xla/xla/debug_options_flags.cc @@ -2815,6 +2815,11 @@ void MakeDebugOptionsFlags(std::vector* flag_list, bool_setter_for(&DebugOptions::set_xla_gpu_mock_custom_calls), debug_options->xla_gpu_mock_custom_calls(), "Replace custom calls with noop operations.")); + flag_list->push_back( + tsl::Flag("xla_cpu_mock_custom_calls", + bool_setter_for(&DebugOptions::set_xla_cpu_mock_custom_calls), + debug_options->xla_cpu_mock_custom_calls(), + "Replace custom calls with noop operations on CPU.")); flag_list->push_back(tsl::Flag( "xla_gpu_enable_while_loop_double_buffering", bool_setter_for( diff --git a/third_party/xla/xla/service/cpu/thunk_emitter.cc b/third_party/xla/xla/service/cpu/thunk_emitter.cc index 10a36e3d5eeeba..13ae4aa5e6d1e7 100644 --- a/third_party/xla/xla/service/cpu/thunk_emitter.cc +++ b/third_party/xla/xla/service/cpu/thunk_emitter.cc @@ -1234,9 +1234,21 @@ absl::StatusOr ThunkEmitter::EmitCustomCallThunk( ABSL_ASSIGN_OR_RETURN(auto op_buffers, GetOpBuffers( instruction, buffer_assignment_)); - return ThunkSequence::Of(ThunkInfo(instruction), - custom_call_target, op_buffers, - backend_config_str, version); + absl::StatusOr> custom_call_thunk = + CustomCallThunk::Create(ThunkInfo(instruction), custom_call_target, + op_buffers, backend_config_str, version); + + if (custom_call_thunk.ok()) { + ThunkSequence thunks; + thunks.push_back(std::move(*custom_call_thunk)); + return thunks; + } + if (hlo_module_config_.debug_options().xla_cpu_mock_custom_calls()) { + // xla_cpu_mock_custom_calls=true means we won't emit thunks for custom + // call targets that couldn't be found. + return ThunkSequence::Empty(); + } + return custom_call_thunk.status(); } absl::StatusOr ThunkEmitter::EmitSliceToDynamicThunk( diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto index c019b460b6fced..31b3003428a1cf 100644 --- a/third_party/xla/xla/xla.proto +++ b/third_party/xla/xla/xla.proto @@ -340,6 +340,9 @@ message DebugOptions { // the flag for more flexible control if necessary. optional string xla_cpu_max_isa = 333; + // Replace custom calls with noop operations on CPU. + optional bool xla_cpu_mock_custom_calls = 548; + // The number of parts to split the LLVM module into before codegen. This // allows XLA to compile all parts in parallel, and resolve kernel symbols // from different dynamic libraries. From 02ed1c36242fe3cebc62bb44df45fa6cf54dd713 Mon Sep 17 00:00:00 2001 From: Saketh Reddy Pingili Date: Mon, 21 Sep 2026 10:02:08 -0700 Subject: [PATCH 12/17] PR #49146: [XLA:CPU] Replace strict-aliasing-violating type punning with absl::bit_cast Imported from GitHub PR https://github.com/openxla/xla/pull/49146 ## [XLA:CPU] Replace strict-aliasing-violating type punning with absl::bit_cast ### Problem In `onednn_ops_rewriter.cc` and `onednn_op_thunk_test.cc`, converting the IEEE-754 bit representation of `float` epsilon to an `int32_t` config field used pointer-dereference type-punning (`*reinterpret_cast(&eps)`). Dereferencing an aliased pointer of an incompatible type violates the C++ strict aliasing rule ([basic.lval]), which is undefined behavior and risks miscompilations under aggressive optimization levels. ### Solution Replace the pointer type-puns with `absl::bit_cast(...)`. This preserves the exact IEEE-754 bit pattern without numeric conversion while adhering strictly to standard aliasing rules. ### Changes - `xla/service/cpu/onednn_ops_rewriter.cc`: Include `absl/base/casts.h` and use `absl::bit_cast(eps)`. - `xla/service/cpu/BUILD`: Add `@com_google_absl//absl/base:casts` to `onednn_ops_rewriter` deps. - `xla/backends/cpu/runtime/onednn/onednn_op_thunk_test.cc`: Include `absl/base/casts.h` and use `absl::bit_cast(epsilon)`. - `xla/backends/cpu/runtime/onednn/BUILD`: Add `@com_google_absl//absl/base:casts` to `onednn_op_thunk_test` deps. ### Testing - `git diff --check`: Passed with no whitespace or syntax formatting issues. - Verified exact 4-file diff matches project include and dependency conventions. ### Risk Low risk. Behavior-preserving change that preserves identical bit-level values while removing undefined behavior. Copybara import of the project: -- 166e8f6bd7c144aa220220b81aa96ffcbd943da8 by saketh reddy pingili : [XLA:CPU] Replace strict-aliasing-violating type punning with absl::bit_cast -- a502f377ac91ce49e01175e0b87489e38ffb3d14 by saketh reddy pingili : Fix absl dependency target to @com_google_absl//absl/base Merging this change closes #49146 PiperOrigin-RevId: 985346019 --- third_party/xla/xla/backends/cpu/runtime/onednn/BUILD | 1 + .../xla/backends/cpu/runtime/onednn/onednn_op_thunk_test.cc | 5 +++-- third_party/xla/xla/service/cpu/BUILD | 1 + third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc | 3 ++- 4 files changed, 7 insertions(+), 3 deletions(-) diff --git a/third_party/xla/xla/backends/cpu/runtime/onednn/BUILD b/third_party/xla/xla/backends/cpu/runtime/onednn/BUILD index f9fd6fd0b94494..d5d78f25d4fb1a 100644 --- a/third_party/xla/xla/backends/cpu/runtime/onednn/BUILD +++ b/third_party/xla/xla/backends/cpu/runtime/onednn/BUILD @@ -185,6 +185,7 @@ onednn_cc_test( "//xla/tsl/platform:env", "//xla/tsl/platform:statusor", "//xla/tsl/platform:test", + "@com_google_absl//absl/base", "@com_google_googletest//:gtest_main", "@eigen_archive//:eigen3", ], diff --git a/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_op_thunk_test.cc b/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_op_thunk_test.cc index f2cbdb93f5c25e..e5943177e7700a 100644 --- a/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_op_thunk_test.cc +++ b/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_op_thunk_test.cc @@ -15,9 +15,10 @@ limitations under the License. #include "xla/backends/cpu/runtime/onednn/onednn_op_thunk.h" +#include #include -// #include "gtest/gtest.h" +#include "absl/base/casts.h" #include "xla/array2d.h" #include "xla/backends/cpu/runtime/buffer_allocations.h" #include "xla/backends/cpu/runtime/thunk.h" @@ -279,7 +280,7 @@ TEST(OneDnnOpThunkTest, SimpleOneDnnLayerNormThunk) { OneDnnNormConfig ln_cfg; ln_cfg.set_rescale(OneDnnNormConfig::SCALE_AND_SHIFT); float epsilon = 1e-5f; - int32_t epsilon_bits = *reinterpret_cast(&epsilon); + int32_t epsilon_bits = absl::bit_cast(epsilon); ln_cfg.set_epsilon_typecast(epsilon_bits); OneDnnOpThunk::OneDnnOpConfig config = ln_cfg; diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD index b89b3e5d863ee5..dbf5d073be1a37 100644 --- a/third_party/xla/xla/service/cpu/BUILD +++ b/third_party/xla/xla/service/cpu/BUILD @@ -1606,6 +1606,7 @@ cc_library( "//xla/service:pattern_matcher", "//xla/tsl/platform:errors", "//xla/tsl/platform:statusor", + "@com_google_absl//absl/base", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/status", diff --git a/third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc b/third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc index 0c172b962681b1..c3c730aa3bf81b 100644 --- a/third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc +++ b/third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include "absl/base/casts.h" #include "absl/container/flat_hash_set.h" #include "absl/container/inlined_vector.h" #include "absl/status/status.h" @@ -546,7 +547,7 @@ class OneDnnOpsRewriterVisitor : public DfsHloRewriteVisitor { OneDnnNormConfig* ln_config = backend_config.mutable_onednn_layer_norm_config(); ln_config->set_rescale(OneDnnNormConfig::SCALE_AND_SHIFT); - ln_config->set_epsilon_typecast(*(reinterpret_cast(&eps))); + ln_config->set_epsilon_typecast(absl::bit_cast(eps)); ABSL_RETURN_IF_ERROR(ln_call->set_backend_config(backend_config)); if (convert_instr != nullptr && is_bf16orfp16_convert && From 2b772f60527938c921ca332b2b30648cba5ff921 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 21 Sep 2026 10:19:43 -0700 Subject: [PATCH 13/17] Exclude TFLite GPU delegate targets from TAP and `tensorflow.v2_gpu_pascal` PiperOrigin-RevId: 985357669 --- tensorflow/lite/delegates/gpu/BUILD | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tensorflow/lite/delegates/gpu/BUILD b/tensorflow/lite/delegates/gpu/BUILD index 05a00673716d7f..f71b2323407861 100644 --- a/tensorflow/lite/delegates/gpu/BUILD +++ b/tensorflow/lite/delegates/gpu/BUILD @@ -389,7 +389,10 @@ cc_library( cc_test( name = "delegate_test", srcs = ["delegate_test.cc"], - tags = tf_gpu_tests_tags() + ["tflite_not_portable_ios"], + tags = tf_gpu_tests_tags() + [ + "notap", + "tflite_not_portable_ios", + ], deps = [ ":delegate", "@com_google_googletest//:gtest_main", From 2eff51d420ccb9d21fd800cdfca0d0d8e1dfef4a Mon Sep 17 00:00:00 2001 From: Alex Pivovarov Date: Mon, 21 Sep 2026 10:23:41 -0700 Subject: [PATCH 14/17] Log profile annotation when creating command buffer thunks. Adds command_buffer_profile_annotation to the VLOG(2) message in the command buffer conversion pass to make it easier to identify which command buffer is being created. Example: Creating command buffer thunk command_buffer_1 with the following thunks: loop_xor_fusion, sort.1, loop_and_convert_fusion PiperOrigin-RevId: 985360175 --- .../xla/backends/gpu/runtime/command_buffer_conversion_pass.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 f6dbb577deb27a..3730af6a1cd1db 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 @@ -572,7 +572,8 @@ ConvertThunksToCommandBuffer( !debug_options.xla_enable_command_buffers_during_profiling()) { thunk_info.profile_annotation += " (disabled for profiling)"; } - VLOG(2) << "Creating command buffer thunk with the following thunks: " + VLOG(2) << "Creating command buffer thunk " + << command_buffer_profile_annotation << " with the following thunks: " << absl::StrJoin( thunks_to_convert, ", ", [](std::string* out, const std::unique_ptr& thunk) { From 62d80d307e41ab82cd2f90013de2b8bf8a59da87 Mon Sep 17 00:00:00 2001 From: "Patrick C. Toulme" Date: Mon, 21 Sep 2026 10:49:21 -0700 Subject: [PATCH 15/17] [XLA:TPU] Build fewer alias analyses in copy insertion Split CopyInsertion::AddSpecialCaseCopies into ComputeSpecialCaseCopies and InsertSpecialCaseCopies. A backend can now compute its own special case copies from the same alias analysis instead of building a second one. The inserted copies are unchanged. PiperOrigin-RevId: 985376755 --- third_party/xla/xla/service/copy_insertion.cc | 42 ++++++++++++------- third_party/xla/xla/service/copy_insertion.h | 16 ++++++- 2 files changed, 43 insertions(+), 15 deletions(-) diff --git a/third_party/xla/xla/service/copy_insertion.cc b/third_party/xla/xla/service/copy_insertion.cc index b9ca9d683cdfa6..bb85d255e64bc5 100644 --- a/third_party/xla/xla/service/copy_insertion.cc +++ b/third_party/xla/xla/service/copy_insertion.cc @@ -1608,7 +1608,19 @@ absl::Status CopyInsertion::AddSpecialCaseCopies( HloModule* module, CustomBufferAnalysisFn custom_buffer_analysis) { ABSL_ASSIGN_OR_RETURN(std::unique_ptr alias_analysis, HloAliasAnalysis::Run(module, alias_info_)); + ABSL_ASSIGN_OR_RETURN( + HloInstructionMap> instructions_to_copy, + ComputeSpecialCaseCopies(call_graph, execution_threads, module, + *alias_analysis, custom_buffer_analysis)); + return InsertSpecialCaseCopies(instructions_to_copy); +} +absl::StatusOr>> +CopyInsertion::ComputeSpecialCaseCopies( + const CallGraph& call_graph, + const absl::flat_hash_set& execution_threads, + HloModule* module, const HloAliasAnalysis& alias_analysis, + CustomBufferAnalysisFn custom_buffer_analysis) { // Identify which shape indices of which instructions need to be copied. Store // these results in 'instructions_to_copy'. HloInstructionMap> instructions_to_copy; @@ -1648,8 +1660,8 @@ absl::Status CopyInsertion::AddSpecialCaseCopies( // Also, locate all input-output aliasing violations for operations that // cannot be done in place. Such aliasing can be created when some copies are // removed too aggressively by CopyRemoval. - for (const HloValue* value : alias_analysis->dataflow_analysis().values()) { - const HloBuffer& buffer = alias_analysis->GetBufferContainingValue(*value); + for (const HloValue* value : alias_analysis.dataflow_analysis().values()) { + const HloBuffer& buffer = alias_analysis.GetBufferContainingValue(*value); if (buffer.values().size() > 1 && ValueIsReadOnly(*value)) { VLOG(2) << "Value " << value->ToShortString() << " is read only, but its buffer contains more than one value. " @@ -1676,13 +1688,12 @@ absl::Status CopyInsertion::AddSpecialCaseCopies( use.operand_number == 0) { continue; } - if (!alias_analysis->dataflow_analysis() - .CanShareOperandBufferWithUser( - /*operand=*/use.instruction->mutable_operand( - use.operand_number), - /*operand_index=*/use.operand_index, - /*user=*/position.instruction, - /*user_index=*/position.index, alias_info_)) { + if (!alias_analysis.dataflow_analysis().CanShareOperandBufferWithUser( + /*operand=*/use.instruction->mutable_operand( + use.operand_number), + /*operand_index=*/use.operand_index, + /*user=*/position.instruction, + /*user_index=*/position.index, alias_info_)) { VLOG(2) << "Adding back copy: " << use.instruction->operand(use.operand_number)->ToString() << "@" << use.operand_index.ToString() @@ -1699,7 +1710,7 @@ absl::Status CopyInsertion::AddSpecialCaseCopies( if (custom_buffer_analysis) { VLOG(2) << "Running custom buffer analysis"; - custom_buffer_analysis(module, *alias_analysis, add_index_to_copy); + custom_buffer_analysis(module, alias_analysis, add_index_to_copy); } // Identify copies which must be added at root instructions @@ -1719,9 +1730,9 @@ absl::Status CopyInsertion::AddSpecialCaseCopies( root->shape(), [&](const Shape& subshape, const ShapeIndex& index) { bool copy_replicated = policy.copy_root_replicated_buffers || - ShouldCopyConditionalRootAt(node, *alias_analysis, index); + ShouldCopyConditionalRootAt(node, alias_analysis, index); std::vector buffers_at_index = - alias_analysis->ComputeBuffersAt(root, index); + alias_analysis.ComputeBuffersAt(root, index); bool buffer_seen_before = false; for (const HloBuffer* buffer : buffers_at_index) { buffer_seen_before |= !seen.emplace(buffer, index).second; @@ -1752,7 +1763,7 @@ absl::Status CopyInsertion::AddSpecialCaseCopies( }); for (const auto& pair : - alias_analysis->dataflow_analysis().GetInstructionValueSet(root)) { + alias_analysis.dataflow_analysis().GetInstructionValueSet(root)) { const ShapeIndex& index = pair.first; const HloValueSet& value_set = pair.second; for (const HloValue* value : value_set.values()) { @@ -1766,8 +1777,11 @@ absl::Status CopyInsertion::AddSpecialCaseCopies( } } } + return instructions_to_copy; +} - // Add copy instructions indicated in 'instructions_to_copy' to the module. +absl::Status CopyInsertion::InsertSpecialCaseCopies( + const HloInstructionMap>& instructions_to_copy) { for (const auto& pair : instructions_to_copy) { HloInstruction* instruction = pair.first; const ShapeTree& indices_to_copy = pair.second; diff --git a/third_party/xla/xla/service/copy_insertion.h b/third_party/xla/xla/service/copy_insertion.h index 135a7b94cde138..31f928a97b84db 100644 --- a/third_party/xla/xla/service/copy_insertion.h +++ b/third_party/xla/xla/service/copy_insertion.h @@ -32,6 +32,7 @@ limitations under the License. #include "xla/hlo/pass/hlo_pass_interface.h" #include "xla/service/call_graph.h" #include "xla/service/hlo_value.h" +#include "xla/shape_tree.h" #include "xla/shape_util.h" namespace xla { @@ -114,12 +115,25 @@ class CopyInsertion : public HloModulePass { CustomBufferAnalysisFn custom_buffer_analysis = nullptr); protected: - // Override which requires the caller to pass in a call graph. + // Override which requires the caller to pass in a call graph. Runs the alias + // analysis, then ComputeSpecialCaseCopies and InsertSpecialCaseCopies. virtual absl::Status AddSpecialCaseCopies( const CallGraph& call_graph, const absl::flat_hash_set& execution_threads, HloModule* module, CustomBufferAnalysisFn custom_buffer_analysis); + // Records the shape indices that need special case copies without inserting + // them. `alias_analysis` must be current for `module`. + absl::StatusOr>> ComputeSpecialCaseCopies( + const CallGraph& call_graph, + const absl::flat_hash_set& execution_threads, + HloModule* module, const HloAliasAnalysis& alias_analysis, + CustomBufferAnalysisFn custom_buffer_analysis); + + // Inserts the copies computed by ComputeSpecialCaseCopies. + static absl::Status InsertSpecialCaseCopies( + const HloInstructionMap>& instructions_to_copy); + // Add copies for conditional instructions. virtual absl::Status AddCopiesForConditional( const HloAliasAnalysis& alias_analysis, HloInstruction* conditional); From 48371417a75e19cbab9804d7b873979a45b35de2 Mon Sep 17 00:00:00 2001 From: Yue Sheng Date: Mon, 21 Sep 2026 10:55:07 -0700 Subject: [PATCH 16/17] [Mosaic TPU] Support 0D convolutions. PiperOrigin-RevId: 985380666 --- third_party/xla/xla/mosaic/dialect/tpu/tpu.td | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu.td b/third_party/xla/xla/mosaic/dialect/tpu/tpu.td index d4e241c0f16f3e..730141fdda7f5d 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu.td +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu.td @@ -67,23 +67,23 @@ def TPU_ConvDimensionNumbersAttr : TPU_Attr<"ConvDimensionNumbers", "conv_dimens let parameters = (ins "int64_t":$input_batch_dimension, "int64_t":$input_feature_dimension, - ArrayRefParameter<"int64_t", "">:$input_spatial_dimensions, + OptionalArrayRefParameter<"int64_t", "">:$input_spatial_dimensions, "int64_t":$kernel_input_feature_dimension, "int64_t":$kernel_output_feature_dimension, - ArrayRefParameter<"int64_t", "">:$kernel_spatial_dimensions, + OptionalArrayRefParameter<"int64_t", "">:$kernel_spatial_dimensions, "int64_t":$output_batch_dimension, "int64_t":$output_feature_dimension, - ArrayRefParameter<"int64_t", "">:$output_spatial_dimensions + OptionalArrayRefParameter<"int64_t", "">:$output_spatial_dimensions ); let assemblyFormat = "`<` $input_batch_dimension `,` " "$input_feature_dimension `,` " - "`[` $input_spatial_dimensions `]` `,` " + "`[` (`]`):($input_spatial_dimensions^ `]`)? `,` " "$kernel_input_feature_dimension `,` " "$kernel_output_feature_dimension `,` " - "`[` $kernel_spatial_dimensions `]` `,` " + "`[` (`]`):($kernel_spatial_dimensions^ `]`)? `,` " "$output_batch_dimension `,` " "$output_feature_dimension `,` " - "`[` $output_spatial_dimensions `]` `>`"; + "`[` (`]`):($output_spatial_dimensions^ `]`)? `>`"; let constBuilderCall = "::mlir::tpu::ConvDimensionNumbersAttr::get($_builder.getContext(), $0)"; } From 02fcca3c8685c1850fdd08e063db7ec897e48c7b Mon Sep 17 00:00:00 2001 From: Ilya Tikhonovskiy Date: Mon, 21 Sep 2026 10:59:50 -0700 Subject: [PATCH 17/17] [XLA:GPU] Skip DotAlgorithmRewriter on pre-Ampere architectures where BF16 does not exist. Pass GpuComputeCapability into DotAlgorithmRewriter so it can check target architecture capabilities. Pre-Ampere GPUs lack native BF16 Tensor Cores, where BF16 matmuls are upcast to F32 for performance. PiperOrigin-RevId: 985383850 --- .../backends/gpu/autotuner/factory_cuda.cc | 3 +- .../backends/gpu/autotuner/factory_rocm.cc | 3 +- .../gpu/autotuner/fission_backend_test.cc | 3 +- .../xla/xla/backends/gpu/transforms/BUILD | 3 + .../gpu/transforms/dot_algorithm_rewriter.cc | 8 + .../gpu/transforms/dot_algorithm_rewriter.h | 7 +- .../transforms/dot_algorithm_rewriter_test.cc | 149 ++++++++++++++++-- .../xla/xla/service/gpu/gpu_compiler.cc | 2 +- third_party/xla/xla/tools/hlo_opt/gpu_opt.cc | 2 +- 9 files changed, 164 insertions(+), 16 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/autotuner/factory_cuda.cc b/third_party/xla/xla/backends/gpu/autotuner/factory_cuda.cc index 8d9094a35cd758..79632575c51fd3 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/factory_cuda.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/factory_cuda.cc @@ -60,7 +60,8 @@ std::unique_ptr GetCublasLtRewriterPipeline( device_description.gpu_compute_capability(), device_description.runtime_version(), GemmRewriterOptions{GemmRewriterOptions::DType::kFp8Only})); - pipeline->AddPass(std::make_unique()); + pipeline->AddPass(std::make_unique( + device_description.gpu_compute_capability())); pipeline->AddPass(std::make_unique( device_description.gpu_compute_capability(), device_description.runtime_version(), diff --git a/third_party/xla/xla/backends/gpu/autotuner/factory_rocm.cc b/third_party/xla/xla/backends/gpu/autotuner/factory_rocm.cc index 6e195de99ff694..e8fbcb7e5e28fa 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/factory_rocm.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/factory_rocm.cc @@ -67,7 +67,8 @@ std::unique_ptr GetGemmRewriterPipeline( device_description.gpu_compute_capability(), device_description.runtime_version(), GemmRewriterOptions{DType::kFp8Only})); - pipeline->AddPass(std::make_unique()); + pipeline->AddPass(std::make_unique( + device_description.gpu_compute_capability())); pipeline->AddPass(std::make_unique( device_description.gpu_compute_capability(), device_description.runtime_version(), diff --git a/third_party/xla/xla/backends/gpu/autotuner/fission_backend_test.cc b/third_party/xla/xla/backends/gpu/autotuner/fission_backend_test.cc index f9dd0c77246872..be8b7dc4a2e7ba 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/fission_backend_test.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/fission_backend_test.cc @@ -144,7 +144,8 @@ std::unique_ptr GetCublasRewriterPipeline( const se::DeviceDescription& device_description) { auto pipeline = std::make_unique("fission_pipeline"); pipeline->AddPass(std::make_unique()); - pipeline->AddPass(std::make_unique()); + pipeline->AddPass(std::make_unique( + device_description.gpu_compute_capability())); for (GemmRewriterOptions::DType dtype : {GemmRewriterOptions::DType::kFp8Only, GemmRewriterOptions::DType::kNonFp8Only}) { diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index f825adfc9e98d7..b1590dd9a34683 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -29,6 +29,7 @@ cc_library( "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", "//xla/service:hlo_creation_utils", + "//xla/stream_executor:device_description", "//xla/tsl/platform:statusor", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log", @@ -49,6 +50,8 @@ xla_cc_test( "//xla/hlo/testlib:filecheck", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/service:hlo_module_config", + "//xla/stream_executor:device_description", + "//xla/stream_executor/cuda:cuda_compute_capability", "@com_google_googletest//:gtest_main", ], ) diff --git a/third_party/xla/xla/backends/gpu/transforms/dot_algorithm_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/dot_algorithm_rewriter.cc index 0459b4c0c8b2fa..64d4a4186a322d 100644 --- a/third_party/xla/xla/backends/gpu/transforms/dot_algorithm_rewriter.cc +++ b/third_party/xla/xla/backends/gpu/transforms/dot_algorithm_rewriter.cc @@ -330,6 +330,14 @@ void RewriteF32ToTF32X3(HloInstruction* instr) { absl::StatusOr DotAlgorithmRewriter::RunImpl( HloModule* module, const absl::flat_hash_set& execution_threads) { + // Pre-Ampere architectures (e.g. Volta SM70, Turing SM75) do not have BF16 + // Tensor Cores or TF32 hardware support. On pre-Ampere GPUs, BF16 matmuls are + // explicitly upcast to F32 for performance, so we do not rewrite them here. + const auto* cuda_cc = gpu_version_.cuda_compute_capability(); + if (cuda_cc != nullptr && !cuda_cc->IsAtLeastAmpere()) { + return false; + } + bool changed = false; bool default_to_bf16 = module->config() .debug_options() diff --git a/third_party/xla/xla/backends/gpu/transforms/dot_algorithm_rewriter.h b/third_party/xla/xla/backends/gpu/transforms/dot_algorithm_rewriter.h index e513b06e645644..90ca88566c0a8d 100644 --- a/third_party/xla/xla/backends/gpu/transforms/dot_algorithm_rewriter.h +++ b/third_party/xla/xla/backends/gpu/transforms/dot_algorithm_rewriter.h @@ -22,6 +22,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/pass/hlo_pass_interface.h" +#include "xla/stream_executor/device_description.h" namespace xla::gpu { @@ -47,7 +48,8 @@ namespace xla::gpu { // emitters like Triton. class DotAlgorithmRewriter : public HloModulePass { public: - DotAlgorithmRewriter() = default; + explicit DotAlgorithmRewriter(se::GpuComputeCapability gpu_version) + : gpu_version_(gpu_version) {} absl::string_view name() const override { return "dot-algorithm-rewriter"; } static absl::StatusOr MakeMultiplyForBF16BF16F32( @@ -67,6 +69,9 @@ class DotAlgorithmRewriter : public HloModulePass { absl::StatusOr RunImpl( HloModule* module, const absl::flat_hash_set& execution_threads) override; + + private: + se::GpuComputeCapability gpu_version_; }; } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/transforms/dot_algorithm_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/dot_algorithm_rewriter_test.cc index 47d4b3e68276ac..86c29f6968d43e 100644 --- a/third_party/xla/xla/backends/gpu/transforms/dot_algorithm_rewriter_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/dot_algorithm_rewriter_test.cc @@ -22,6 +22,8 @@ limitations under the License. #include "xla/hlo/testlib/filecheck.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/service/hlo_module_config.h" +#include "xla/stream_executor/cuda/cuda_compute_capability.h" +#include "xla/stream_executor/device_description.h" #include "xla/xla.pb.h" #include "xla/xla_data.pb.h" @@ -50,8 +52,11 @@ TEST_F(DotAlgorithmRewriterTest, DefaultToBF16) { ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text, config)); - ASSERT_OK_AND_ASSIGN(auto pass_result, - RunHloPass(DotAlgorithmRewriter(), module.get())); + ASSERT_OK_AND_ASSIGN( + auto pass_result, + RunHloPass(DotAlgorithmRewriter( + stream_executor::CudaComputeCapability::Hopper()), + module.get())); EXPECT_TRUE(pass_result); const char* expected = R"( @@ -88,8 +93,11 @@ TEST_F(DotAlgorithmRewriterTest, NoDefaultToBF16) { ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text, config)); - ASSERT_OK_AND_ASSIGN(bool pass_result, - RunHloPass(DotAlgorithmRewriter(), module.get())); + ASSERT_OK_AND_ASSIGN( + bool pass_result, + RunHloPass(DotAlgorithmRewriter( + stream_executor::CudaComputeCapability::Hopper()), + module.get())); EXPECT_FALSE(pass_result); } @@ -113,8 +121,11 @@ TEST_F(DotAlgorithmRewriterTest, DefaultToBF16_NonF32Result) { ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text, config)); - ASSERT_OK_AND_ASSIGN(auto pass_result, - RunHloPass(DotAlgorithmRewriter(), module.get())); + ASSERT_OK_AND_ASSIGN( + auto pass_result, + RunHloPass(DotAlgorithmRewriter( + stream_executor::CudaComputeCapability::Hopper()), + module.get())); EXPECT_FALSE(pass_result); } @@ -138,8 +149,11 @@ TEST_F(DotAlgorithmRewriterTest, DefaultToBF16_NonF32Operands) { ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text, config)); - ASSERT_OK_AND_ASSIGN(auto pass_result, - RunHloPass(DotAlgorithmRewriter(), module.get())); + ASSERT_OK_AND_ASSIGN( + auto pass_result, + RunHloPass(DotAlgorithmRewriter( + stream_executor::CudaComputeCapability::Hopper()), + module.get())); EXPECT_FALSE(pass_result); } @@ -164,10 +178,125 @@ TEST_F(DotAlgorithmRewriterTest, DefaultToBF16_HighestPrecision) { ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text, config)); - ASSERT_OK_AND_ASSIGN(auto pass_result, - RunHloPass(DotAlgorithmRewriter(), module.get())); + ASSERT_OK_AND_ASSIGN( + auto pass_result, + RunHloPass(DotAlgorithmRewriter( + stream_executor::CudaComputeCapability::Hopper()), + module.get())); EXPECT_FALSE(pass_result); } +TEST_F(DotAlgorithmRewriterTest, SkipOnP100) { + const char* hlo_text = R"hlo( + HloModule test + + ENTRY test { + p0 = f32[32,32] parameter(0) + p1 = f32[32,32] parameter(1) + ROOT dot = f32[32,32] dot(p0, p1), + lhs_contracting_dims={1}, + rhs_contracting_dims={0} + } + )hlo"; + + HloModuleConfig config = GetModuleConfigForTest(); + DebugOptions debug_options = config.debug_options(); + debug_options.set_xla_gpu_default_to_alg_dot_bf16_bf16_f32(true); + config.set_debug_options(debug_options); + + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_text, config)); + ASSERT_OK_AND_ASSIGN( + auto pass_result, + RunHloPass(DotAlgorithmRewriter( + stream_executor::CudaComputeCapability::Pascal()), + module.get())); + EXPECT_FALSE(pass_result); +} + +TEST_F(DotAlgorithmRewriterTest, SkipOnV100) { + const char* hlo_text = R"hlo( + HloModule test + + ENTRY test { + p0 = f32[32,32] parameter(0) + p1 = f32[32,32] parameter(1) + ROOT dot = f32[32,32] dot(p0, p1), + lhs_contracting_dims={1}, + rhs_contracting_dims={0} + } + )hlo"; + + HloModuleConfig config = GetModuleConfigForTest(); + DebugOptions debug_options = config.debug_options(); + debug_options.set_xla_gpu_default_to_alg_dot_bf16_bf16_f32(true); + config.set_debug_options(debug_options); + + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_text, config)); + ASSERT_OK_AND_ASSIGN( + auto pass_result, + RunHloPass( + DotAlgorithmRewriter(stream_executor::CudaComputeCapability::Volta()), + module.get())); + EXPECT_FALSE(pass_result); +} + +TEST_F(DotAlgorithmRewriterTest, SkipOnTuring) { + const char* hlo_text = R"hlo( + HloModule test + + ENTRY test { + p0 = f32[32,32] parameter(0) + p1 = f32[32,32] parameter(1) + ROOT dot = f32[32,32] dot(p0, p1), + lhs_contracting_dims={1}, + rhs_contracting_dims={0} + } + )hlo"; + + HloModuleConfig config = GetModuleConfigForTest(); + DebugOptions debug_options = config.debug_options(); + debug_options.set_xla_gpu_default_to_alg_dot_bf16_bf16_f32(true); + config.set_debug_options(debug_options); + + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_text, config)); + ASSERT_OK_AND_ASSIGN( + auto pass_result, + RunHloPass( + DotAlgorithmRewriter(stream_executor::CudaComputeCapability{7, 5}), + module.get())); + EXPECT_FALSE(pass_result); +} + +TEST_F(DotAlgorithmRewriterTest, RunOnA100) { + const char* hlo_text = R"hlo( + HloModule test + + ENTRY test { + p0 = f32[32,32] parameter(0) + p1 = f32[32,32] parameter(1) + ROOT dot = f32[32,32] dot(p0, p1), + lhs_contracting_dims={1}, + rhs_contracting_dims={0} + } + )hlo"; + + HloModuleConfig config = GetModuleConfigForTest(); + DebugOptions debug_options = config.debug_options(); + debug_options.set_xla_gpu_default_to_alg_dot_bf16_bf16_f32(true); + config.set_debug_options(debug_options); + + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_text, config)); + ASSERT_OK_AND_ASSIGN( + auto pass_result, + RunHloPass(DotAlgorithmRewriter( + stream_executor::CudaComputeCapability::Ampere()), + module.get())); + EXPECT_TRUE(pass_result); +} + } // namespace } // namespace xla::gpu diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc index 69c73b3f4c3f43..24451d43131c8c 100644 --- a/third_party/xla/xla/service/gpu/gpu_compiler.cc +++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc @@ -2042,7 +2042,7 @@ void AddGemmRewriterPasses(HloPassPipeline& pipeline, // I.e. transform single dot into a chain of dots with the default algorithm // that cublas can handle. These dots were inlined by the CallInliner pass // above. - pipeline.AddPass(); + pipeline.AddPass(gpu_version); pipeline.AddPass( gpu_version, toolkit_version, diff --git a/third_party/xla/xla/tools/hlo_opt/gpu_opt.cc b/third_party/xla/xla/tools/hlo_opt/gpu_opt.cc index 18aff86f6bc77f..c3e4bfccbe8c93 100644 --- a/third_party/xla/xla/tools/hlo_opt/gpu_opt.cc +++ b/third_party/xla/xla/tools/hlo_opt/gpu_opt.cc @@ -173,7 +173,7 @@ class GpuOptProvider : public CompiledOptProvider { RegisterPass(alias_info_.get()); RegisterPass(); RegisterPass(); - RegisterPass(); + RegisterPass(gpu_compute_capability); RegisterPass(); RegisterPass(); RegisterPass();