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) 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", 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/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", 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", ) 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/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/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/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/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/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) { 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/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)); } 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", + ], +) 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/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)"; } 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 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)); 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); 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 && 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/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}} 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(); 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(); 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; } 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.