diff --git a/.bazelrc b/.bazelrc index fa9874a91a2302..c41207afe93d54 100644 --- a/.bazelrc +++ b/.bazelrc @@ -293,13 +293,13 @@ common:mkl_aarch64 --config=mkl_aarch64_threadpool common:cuda12_version --repo_env=HERMETIC_CUDA_VERSION="12.5.1" common:cuda12_version --repo_env=HERMETIC_CUDNN_VERSION="9.10.0" common:cuda12_version --repo_env=HERMETIC_NVSHMEM_VERSION="3.2.5" -common:cuda12_version --repo_env=HERMETIC_NCCL_VERSION="2.29.7" +common:cuda12_version --repo_env=HERMETIC_NCCL_VERSION="2.30.7" common:cuda12_version --repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_60,sm_70,sm_80,sm_89,compute_90" common:cuda13_version --repo_env=HERMETIC_CUDA_VERSION="13.0.0" common:cuda13_version --repo_env=HERMETIC_CUDNN_VERSION="9.12.0" common:cuda13_version --repo_env=HERMETIC_NVSHMEM_VERSION="3.3.20" -common:cuda13_version --repo_env=HERMETIC_NCCL_VERSION="2.29.7" +common:cuda13_version --repo_env=HERMETIC_NCCL_VERSION="2.30.7" common:cuda13_version --repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_75,sm_80,sm_90,sm_100,compute_120" common:cuda_version --config=cuda12_version diff --git a/MODULE.bazel b/MODULE.bazel index 6d285fee56ca40..f3177c24804133 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -86,9 +86,9 @@ bazel_dep(name = "libprotobuf-mutator", version = "1.3", repo_name = "libprotobu bazel_dep(name = "rules_ml_toolchain") archive_override( module_name = "rules_ml_toolchain", - integrity = "sha256-OwVoeEJCcEHGXRvC9K7aOjB5VXEg9b6MNGkAh6iMXeU=", - strip_prefix = "rules_ml_toolchain-2eddbc595cc0bbe650c2640204f66b14f015f1a8", - urls = ["https://github.com/google-ml-infra/rules_ml_toolchain/archive/2eddbc595cc0bbe650c2640204f66b14f015f1a8.tar.gz"], + integrity = "sha256-izPVqzV2I7ox7EYzRzcRY6ogeLcHUr4LRP2xMrkWh6A=", + strip_prefix = "rules_ml_toolchain-236a498d8624bef88e35ce2518a833204b0db19f", + urls = ["https://github.com/google-ml-infra/rules_ml_toolchain/archive/236a498d8624bef88e35ce2518a833204b0db19f.tar.gz"], ) bazel_dep(name = "xla", repo_name = "xla") diff --git a/RELEASE.md b/RELEASE.md index 575ac36ab4b3b6..1c82b90434a66d 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -40,6 +40,12 @@ In `tensorflow/c/experimental/filesystem/filesystem_interface.h`, removed `TF_Tr `tf.image.adjust_contrast` can now be differentiated with `GradientTape`. Fixes [#126083](https://github.com/tensorflow/tensorflow/issues/126083). +* `tf.nn.softsign` + + * Fixes second-order gradients of `tf.nn.softsign`. Differentiating twice + previously failed with a lookup error because the `SoftsignGrad` + backward op had no registered Python gradient. + * `tf.experimental.numpy` diff --git a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt index f98cc872f05cd7..137bc9e338a4e4 100644 --- a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt +++ b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt @@ -519,9 +519,9 @@ nvidia-cusparse-cu12==12.5.1.3 \ # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.29.7 \ - --hash=sha256:0cf032ee22b560447daf0456108a75e32bd74a4de6c6b64725637a359fa48cd8 \ - --hash=sha256:ecd0a012051abc20c1aa87328841efa8cade3ced65803046e38c2f03c0891fea +nvidia-nccl-cu12==2.30.7 \ + --hash=sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1 \ + --hash=sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt diff --git a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt index 9ddfca552b7545..8b37e5bd8b7bde 100644 --- a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt +++ b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt @@ -519,9 +519,9 @@ nvidia-cusparse-cu12==12.5.1.3 \ # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.29.7 \ - --hash=sha256:0cf032ee22b560447daf0456108a75e32bd74a4de6c6b64725637a359fa48cd8 \ - --hash=sha256:ecd0a012051abc20c1aa87328841efa8cade3ced65803046e38c2f03c0891fea +nvidia-nccl-cu12==2.30.7 \ + --hash=sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1 \ + --hash=sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt diff --git a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt index 15e097176cdbd4..7c0230a0cbe19b 100644 --- a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt +++ b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt @@ -519,9 +519,9 @@ nvidia-cusparse-cu12==12.5.1.3 \ # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.29.7 \ - --hash=sha256:0cf032ee22b560447daf0456108a75e32bd74a4de6c6b64725637a359fa48cd8 \ - --hash=sha256:ecd0a012051abc20c1aa87328841efa8cade3ced65803046e38c2f03c0891fea +nvidia-nccl-cu12==2.30.7 \ + --hash=sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1 \ + --hash=sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt diff --git a/ci/official/requirements_updater/nvidia-requirements.txt b/ci/official/requirements_updater/nvidia-requirements.txt index b979a97190a576..160e5f190a590b 100644 --- a/ci/official/requirements_updater/nvidia-requirements.txt +++ b/ci/official/requirements_updater/nvidia-requirements.txt @@ -12,11 +12,11 @@ nvidia-cufft-cu12>=11.2.3.61,<12.0 nvidia-curand-cu12>=10.3.6.82,<11.0 nvidia-cusolver-cu12>=11.6.3.83,<12.0 nvidia-cusparse-cu12>=12.5.1.3,<13.0 -nvidia-nccl-cu12>=2.29.7,<3.0 +nvidia-nccl-cu12>=2.30.7,<3.0 nvidia-nvjitlink-cu12>=12.5.82,<13.0 nvidia-nvshmem-cu12>=3.2.5 # CUDA 13 wheels -nvidia-nccl-cu13>=2.29.7 +nvidia-nccl-cu13>=2.30.7 nvidia-nvshmem-cu13>=3.3.20 nvidia-cublas>=13.0.0.19 nvidia-cuda-cupti>=13.0.48 diff --git a/ci/official/utilities/code_check_full.bats b/ci/official/utilities/code_check_full.bats index 1f6e6f810f91f7..b01432c7de0949 100644 --- a/ci/official/utilities/code_check_full.bats +++ b/ci/official/utilities/code_check_full.bats @@ -223,7 +223,7 @@ EOF --repo_env=HERMETIC_CUDA_VERSION="12.5.1" \ --repo_env=HERMETIC_CUDNN_VERSION="9.10.0" \ --repo_env=HERMETIC_NVSHMEM_VERSION="3.2.5" \ - --repo_env=HERMETIC_NCCL_VERSION="2.29.7" \ + --repo_env=HERMETIC_NCCL_VERSION="2.30.7" \ "somepath(//tensorflow/tools/pip_package:wheel, " \ "@local_config_cuda//cuda:cudart + "\ "@local_config_cuda//cuda:cudart + "\ @@ -249,7 +249,7 @@ EOF --repo_env=HERMETIC_CUDA_VERSION="12.5.1" \ --repo_env=HERMETIC_CUDNN_VERSION="9.10.0" \ --repo_env=HERMETIC_NVSHMEM_VERSION="3.2.5" \ - --repo_env=HERMETIC_NCCL_VERSION="2.29.7" \ + --repo_env=HERMETIC_NCCL_VERSION="2.30.7" \ --define framework_shared_object=false \ "somepath(//tensorflow/tools/pip_package:wheel, " \ "@local_config_cuda//cuda:cudart + "\ diff --git a/requirements_lock_3_10.txt b/requirements_lock_3_10.txt index 11f87a3b5687e2..e2298fa546a5bb 100644 --- a/requirements_lock_3_10.txt +++ b/requirements_lock_3_10.txt @@ -610,18 +610,16 @@ nvidia-cusparse-cu12==12.5.1.3 \ # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.29.7 \ - --hash=sha256:0cf032ee22b560447daf0456108a75e32bd74a4de6c6b64725637a359fa48cd8 \ - --hash=sha256:ecd0a012051abc20c1aa87328841efa8cade3ced65803046e38c2f03c0891fea +nvidia-nccl-cu12==2.30.7 \ + --hash=sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1 \ + --hash=sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt -nvidia-nccl-cu13==2.29.7 \ - --hash=sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d \ - --hash=sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d \ - --hash=sha256:01c873ba1626b54caa12272ed228dc5b2781545e0ae8ba3f432a8ef1c6d78643 \ - --hash=sha256:e4553a30f34195f3fa1da02a6da3d6337d28f2003943aa0a3d247bbc25fefc42 +nvidia-nccl-cu13==2.30.7 \ + --hash=sha256:ca786ffa5a647c75d4d1f5cc72a6c4f537947e2ba8823d7c8aaf768e7a7b9f77 \ + --hash=sha256:cefa7fdb9710efd0f39c5f1be1d61ff6fc9a996c451265bd7fbdcf9455ed4b50 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt diff --git a/requirements_lock_3_11.txt b/requirements_lock_3_11.txt index f1f748323495a2..df484dd450e71a 100644 --- a/requirements_lock_3_11.txt +++ b/requirements_lock_3_11.txt @@ -610,15 +610,16 @@ nvidia-cusparse-cu12==12.5.1.3 \ # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.29.7 \ - --hash=sha256:0cf032ee22b560447daf0456108a75e32bd74a4de6c6b64725637a359fa48cd8 \ - --hash=sha256:ecd0a012051abc20c1aa87328841efa8cade3ced65803046e38c2f03c0891fea +nvidia-nccl-cu12==2.30.7 \ + --hash=sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1 \ + --hash=sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt -nvidia-nccl-cu13==2.29.7 \ - --hash=sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d \ +nvidia-nccl-cu13==2.30.7 \ + --hash=sha256:ca786ffa5a647c75d4d1f5cc72a6c4f537947e2ba8823d7c8aaf768e7a7b9f77 \ + --hash=sha256:cefa7fdb9710efd0f39c5f1be1d61ff6fc9a996c451265bd7fbdcf9455ed4b50 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt diff --git a/requirements_lock_3_12.txt b/requirements_lock_3_12.txt index a412ddc80dbdb9..af658e13cac1e0 100644 --- a/requirements_lock_3_12.txt +++ b/requirements_lock_3_12.txt @@ -610,17 +610,16 @@ nvidia-cusparse-cu12==12.5.1.3 \ # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.29.7 \ - --hash=sha256:0cf032ee22b560447daf0456108a75e32bd74a4de6c6b64725637a359fa48cd8 \ - --hash=sha256:ecd0a012051abc20c1aa87328841efa8cade3ced65803046e38c2f03c0891fea +nvidia-nccl-cu12==2.30.7 \ + --hash=sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1 \ + --hash=sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt -nvidia-nccl-cu13==2.29.7 \ - --hash=sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d \ - --hash=sha256:01c873ba1626b54caa12272ed228dc5b2781545e0ae8ba3f432a8ef1c6d78643 \ - --hash=sha256:e4553a30f34195f3fa1da02a6da3d6337d28f2003943aa0a3d247bbc25fefc42 +nvidia-nccl-cu13==2.30.7 \ + --hash=sha256:ca786ffa5a647c75d4d1f5cc72a6c4f537947e2ba8823d7c8aaf768e7a7b9f77 \ + --hash=sha256:cefa7fdb9710efd0f39c5f1be1d61ff6fc9a996c451265bd7fbdcf9455ed4b50 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt diff --git a/requirements_lock_3_13.txt b/requirements_lock_3_13.txt index f860aee8c0739a..da604d3fa08573 100644 --- a/requirements_lock_3_13.txt +++ b/requirements_lock_3_13.txt @@ -610,17 +610,16 @@ nvidia-cusparse-cu12==12.5.1.3 \ # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.29.7 \ - --hash=sha256:0cf032ee22b560447daf0456108a75e32bd74a4de6c6b64725637a359fa48cd8 \ - --hash=sha256:ecd0a012051abc20c1aa87328841efa8cade3ced65803046e38c2f03c0891fea +nvidia-nccl-cu12==2.30.7 \ + --hash=sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1 \ + --hash=sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt -nvidia-nccl-cu13==2.29.7 \ - --hash=sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d \ - --hash=sha256:01c873ba1626b54caa12272ed228dc5b2781545e0ae8ba3f432a8ef1c6d78643 \ - --hash=sha256:e4553a30f34195f3fa1da02a6da3d6337d28f2003943aa0a3d247bbc25fefc42 +nvidia-nccl-cu13==2.30.7 \ + --hash=sha256:ca786ffa5a647c75d4d1f5cc72a6c4f537947e2ba8823d7c8aaf768e7a7b9f77 \ + --hash=sha256:cefa7fdb9710efd0f39c5f1be1d61ff6fc9a996c451265bd7fbdcf9455ed4b50 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt diff --git a/requirements_lock_3_14.txt b/requirements_lock_3_14.txt index 5bc65a2d12c992..0a4539230b6f6c 100644 --- a/requirements_lock_3_14.txt +++ b/requirements_lock_3_14.txt @@ -716,17 +716,16 @@ nvidia-cusparse-cu12==12.5.1.3 \ # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.29.7 \ - --hash=sha256:0cf032ee22b560447daf0456108a75e32bd74a4de6c6b64725637a359fa48cd8 \ - --hash=sha256:ecd0a012051abc20c1aa87328841efa8cade3ced65803046e38c2f03c0891fea +nvidia-nccl-cu12==2.30.7 \ + --hash=sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1 \ + --hash=sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt -nvidia-nccl-cu13==2.29.7 \ - --hash=sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d \ - --hash=sha256:01c873ba1626b54caa12272ed228dc5b2781545e0ae8ba3f432a8ef1c6d78643 \ - --hash=sha256:e4553a30f34195f3fa1da02a6da3d6337d28f2003943aa0a3d247bbc25fefc42 +nvidia-nccl-cu13==2.30.7 \ + --hash=sha256:ca786ffa5a647c75d4d1f5cc72a6c4f537947e2ba8823d7c8aaf768e7a7b9f77 \ + --hash=sha256:cefa7fdb9710efd0f39c5f1be1d61ff6fc9a996c451265bd7fbdcf9455ed4b50 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt diff --git a/tensorflow/cc/saved_model/BUILD b/tensorflow/cc/saved_model/BUILD index f604a97fe95c8d..93b4bf515cbde8 100644 --- a/tensorflow/cc/saved_model/BUILD +++ b/tensorflow/cc/saved_model/BUILD @@ -518,18 +518,8 @@ cc_library( "//tensorflow/python:__pkg__", ], deps = [ - ":constants", "//tensorflow/core/protobuf:for_core_protos_cc", - "@com_google_absl//absl/status", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings", - "@tsl//tsl/platform:errors", - "@tsl//tsl/platform:statusor", - ] + if_not_mobile([ - "//tensorflow/core:lib", - ]) + if_android([ - "//tensorflow/core:portable_tensorflow_lib_lite", - ]), + ], alwayslink = True, ) @@ -539,14 +529,9 @@ cc_library( visibility = ["//visibility:public"], deps = if_static([ ":singleprint_impl", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings", + ]) + [ "//tensorflow/core/protobuf:for_core_protos_cc", - ]) + if_not_mobile([ - "//tensorflow/core:lib", - ]) + if_android([ - "//tensorflow/core:portable_tensorflow_lib_lite", - ]), + ], ) tf_cc_test( @@ -555,7 +540,7 @@ tf_cc_test( srcs = ["singleprint_test.cc"], deps = [ ":singleprint", - "//tensorflow/core:protos_all_cc", + "//tensorflow/core/protobuf:for_core_protos_cc", "@com_google_googletest//:gtest_main", ], ) diff --git a/tensorflow/cc/saved_model/fingerprinting.cc b/tensorflow/cc/saved_model/fingerprinting.cc index 9f7519130e6add..b2622a32cb0847 100644 --- a/tensorflow/cc/saved_model/fingerprinting.cc +++ b/tensorflow/cc/saved_model/fingerprinting.cc @@ -27,6 +27,7 @@ limitations under the License. #include "absl/strings/strip.h" #include "tensorflow/cc/saved_model/constants.h" #include "tensorflow/cc/saved_model/fingerprinting_x_platform_utils.h" +#include "tensorflow/cc/saved_model/singleprint.h" #include "tensorflow/core/framework/versions.pb.h" #include "tensorflow/core/graph/regularization/simple_delete.h" #include "tensorflow/core/graph/regularization/util.h" @@ -259,4 +260,10 @@ absl::StatusOr ReadSavedModelFingerprint( return fingerprint_proto; } +absl::StatusOr Singleprint(absl::string_view export_dir) { + TF_ASSIGN_OR_RETURN(FingerprintDef fingerprint_def, + ReadSavedModelFingerprint(export_dir)); + return Singleprint(fingerprint_def); +} + } // namespace tensorflow::saved_model::fingerprinting diff --git a/tensorflow/cc/saved_model/fingerprinting.h b/tensorflow/cc/saved_model/fingerprinting.h index 25adbbe23799da..a4421874d3d2df 100644 --- a/tensorflow/cc/saved_model/fingerprinting.h +++ b/tensorflow/cc/saved_model/fingerprinting.h @@ -16,6 +16,8 @@ limitations under the License. #ifndef TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_H_ #define TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_H_ +#include + #include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "tensorflow/cc/saved_model/singleprint.h" // IWYU pragma: keep. Used for Singleprint(). @@ -33,6 +35,9 @@ absl::StatusOr CreateFingerprintDef( absl::StatusOr ReadSavedModelFingerprint( absl::string_view export_dir); +// Canonical fingerprinting ID for a SavedModel loaded from `export_dir`. +absl::StatusOr Singleprint(absl::string_view export_dir); + } // namespace tensorflow::saved_model::fingerprinting #endif // TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_H_ diff --git a/tensorflow/cc/saved_model/singleprint.cc b/tensorflow/cc/saved_model/singleprint.cc index 4ffc3f533c90f5..a189c8a0293f45 100644 --- a/tensorflow/cc/saved_model/singleprint.cc +++ b/tensorflow/cc/saved_model/singleprint.cc @@ -18,35 +18,10 @@ limitations under the License. #include #include -#include "absl/status/status.h" -#include "absl/status/statusor.h" -#include "absl/strings/string_view.h" -#include "tensorflow/cc/saved_model/constants.h" -#include "tensorflow/core/platform/env.h" -#include "tensorflow/core/platform/path.h" #include "tensorflow/core/protobuf/fingerprint.pb.h" -#include "tsl/platform/statusor.h" namespace tensorflow::saved_model::fingerprinting { -namespace { - -absl::StatusOr ReadSavedModelFingerprint( - absl::string_view export_dir) { - const std::string fingerprint_pb_path = - io::JoinPath(export_dir, kFingerprintFilenamePb); - TF_RETURN_IF_ERROR(Env::Default()->FileExists(fingerprint_pb_path)); - - FingerprintDef fingerprint_proto; - absl::Status result = - ReadBinaryProto(Env::Default(), fingerprint_pb_path, &fingerprint_proto); - if (!result.ok()) return result; - - return fingerprint_proto; -} - -} // namespace - std::string Singleprint(uint64_t graph_def_program_hash, uint64_t signature_def_hash, uint64_t saved_object_graph_hash, @@ -63,10 +38,4 @@ std::string Singleprint(const FingerprintDef& fingerprint) { fingerprint.saved_object_graph_hash(), fingerprint.checkpoint_hash()); } -absl::StatusOr Singleprint(absl::string_view export_dir) { - TF_ASSIGN_OR_RETURN(FingerprintDef fingerprint_def, - ReadSavedModelFingerprint(export_dir)); - return Singleprint(fingerprint_def); -} - } // namespace tensorflow::saved_model::fingerprinting diff --git a/tensorflow/cc/saved_model/singleprint.h b/tensorflow/cc/saved_model/singleprint.h index 95e27d4f6de655..a83de66152ca3a 100644 --- a/tensorflow/cc/saved_model/singleprint.h +++ b/tensorflow/cc/saved_model/singleprint.h @@ -19,8 +19,6 @@ limitations under the License. #include #include -#include "absl/status/statusor.h" -#include "absl/strings/string_view.h" #include "tensorflow/core/protobuf/fingerprint.pb.h" namespace tensorflow::saved_model::fingerprinting { @@ -31,7 +29,6 @@ std::string Singleprint(uint64_t graph_def_program_hash, uint64_t saved_object_graph_hash, uint64_t checkpoint_hash); std::string Singleprint(const FingerprintDef& fingerprint); -absl::StatusOr Singleprint(absl::string_view export_dir); } // namespace tensorflow::saved_model::fingerprinting diff --git a/tensorflow/lite/kernels/slice.cc b/tensorflow/lite/kernels/slice.cc index 597c14f585d25c..86dca38bb7c427 100644 --- a/tensorflow/lite/kernels/slice.cc +++ b/tensorflow/lite/kernels/slice.cc @@ -145,16 +145,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ENSURE_EQ(context, NumElements(begin), NumElements(size)); // If the shape of output is fully specified then resize even if // the input shape is not staticly defined. - // - // A fully specified output shape does not imply the slice is in bounds when - // the input extent is only known at run time. Taking this path leaves the - // output static, so Eval() -- which only re-runs ResizeOutputShape() for a - // dynamic output -- never reaches CalculateOutputShapeVector(), the one place - // `begin` and `size` are checked against the input. Fall through when the - // input has an unspecified dimension so the output is marked dynamic and the - // bounds are validated against the actual extent on every invocation. - if (ShapeHasRank(output->dims) && !HasUnspecifiedDimension(output) && - ShapeHasRank(input->dims) && !HasUnspecifiedDimension(input)) { + if (!HasUnspecifiedDimension(output) && ShapeHasRank(output->dims)) { return kTfLiteOk; } // Postpone allocation of output if any of the indexing tensors is not diff --git a/tensorflow/lite/kernels/slice_test.cc b/tensorflow/lite/kernels/slice_test.cc index e3b55f6116e32c..09b68642fff848 100644 --- a/tensorflow/lite/kernels/slice_test.cc +++ b/tensorflow/lite/kernels/slice_test.cc @@ -102,73 +102,6 @@ class SliceOpModel : public SingleOpModel { class SliceOpTest : public ::testing::TestWithParam {}; -// Model with a dynamic input dimension and a statically shaped output. The -// suite name is deliberately distinct from SliceOpTest, which is a TEST_P -// fixture -- gtest rejects a suite that mixes TEST and TEST_P. -class DynamicInputSliceOpModel : public SingleOpModel { - public: - DynamicInputSliceOpModel(TensorData input_data, - std::initializer_list begin_shape, - std::initializer_list begin_data, - std::initializer_list size_shape, - std::initializer_list size_data, - TensorData output_data) { - input_ = AddInput(input_data); - begin_ = AddConstInput(TensorType_INT32, begin_data, begin_shape); - size_ = AddConstInput(TensorType_INT32, size_data, size_shape); - output_ = AddOutput(output_data); - SetBuiltinOp(BuiltinOperator_SLICE, BuiltinOptions_SliceOptions, - CreateSliceOptions(builder_).Union()); - // Delegates are bypassed: a delegate that claims the SLICE node would set - // the output allocation type itself, so the assertions below would no - // longer describe the built-in CPU kernel. - BuildInterpreter({input_data.shape, begin_shape, size_shape}, - /*num_threads=*/-1, /*allow_fp32_relax_to_fp16=*/false, - /*apply_delegate=*/false); - } - - void SetInput(std::initializer_list data) { - PopulateTensor(input_, data); - } - std::vector GetOutput() { return ExtractVector(output_); } - std::vector GetOutputShape() { return GetTensorShape(output_); } - const TfLiteTensor* GetOutputTensor() { - return interpreter_->tensor(output_); - } - - private: - int input_; - int begin_; - int size_; - int output_; -}; - -// A dynamic input dimension must force the output dynamic so the bounds are -// re-validated in Eval(), even though the declared output shape is static. -TEST(SliceOpDynamicInputTest, DynamicInputStaticOutputValid) { - TensorData input_data(TensorType_FLOAT32, {1, 8}); - input_data.shape_signature = {1, -1}; - TensorData output_data(TensorType_FLOAT32, {1, 4}); - - DynamicInputSliceOpModel m(input_data, {2}, {0, 2}, {2}, {1, 4}, output_data); - m.SetInput({1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0}); - ASSERT_EQ(m.Invoke(), kTfLiteOk); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4})); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({3.0, 4.0, 5.0, 6.0})); - EXPECT_EQ(m.GetOutputTensor()->allocation_type, kTfLiteDynamic); -} - -// The same path must reject a window that does not fit the actual extent. -TEST(SliceOpDynamicInputTest, DynamicInputStaticOutputOutOfBounds) { - TensorData input_data(TensorType_FLOAT32, {1, 4}); - input_data.shape_signature = {1, -1}; - TensorData output_data(TensorType_FLOAT32, {1, 8}); - - DynamicInputSliceOpModel m(input_data, {2}, {0, 0}, {2}, {1, 8}, output_data); - m.SetInput({1.0, 2.0, 3.0, 4.0}); - EXPECT_EQ(m.Invoke(), kTfLiteError); -} - TEST_P(SliceOpTest, In1D) { SliceOpModel m({4}, {1}, {1}, {1}, {2}, TensorType_INT32, TensorType_FLOAT32, GetParam()); diff --git a/tensorflow/python/eager/pywrap_gradient_exclusions.cc b/tensorflow/python/eager/pywrap_gradient_exclusions.cc index 8c34921427963b..6e39edcad17aa1 100644 --- a/tensorflow/python/eager/pywrap_gradient_exclusions.cc +++ b/tensorflow/python/eager/pywrap_gradient_exclusions.cc @@ -429,7 +429,7 @@ absl::optional> OpGradientUnusedInputIndices( absl::optional> OpGradientUnusedOutputIndices( const tensorflow::string &op_name) { - static std::array a = {{ + static std::array a = {{ {"Abs"}, {"AccumulateNV2"}, {"Acos"}, @@ -773,6 +773,7 @@ absl::optional> OpGradientUnusedOutputIndices( {"Softplus"}, {"SoftplusGrad"}, {"Softsign"}, + {"SoftsignGrad"}, {"SpaceToBatch"}, {"SpaceToBatchND"}, {"SpaceToDepth"}, diff --git a/tensorflow/python/kernel_tests/nn_ops/softsign_op_test.py b/tensorflow/python/kernel_tests/nn_ops/softsign_op_test.py index 585c6435aae524..ebb93f657d1fe5 100644 --- a/tensorflow/python/kernel_tests/nn_ops/softsign_op_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/softsign_op_test.py @@ -16,6 +16,7 @@ import numpy as np +from tensorflow.python.eager import backprop from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import test_util @@ -25,6 +26,10 @@ from tensorflow.python.platform import test +def _softsign_grad_grad(activation): + return -2 * np.sign(activation) / (1 + np.abs(activation)) ** 3 + + class SoftsignTest(test.TestCase): def _npSoftsign(self, np_features): @@ -82,6 +87,26 @@ def testNoInts(self): "'features' has DataType int32 not in list of allowed values"): nn_ops.softsign(constant_op.constant(7)).eval() + def testGradGrad(self): + with self.cached_session(): + + def f(x): + with backprop.GradientTape(persistent=True) as tape: + tape.watch(x) + y = nn_ops.softsign(x) + dy = tape.gradient(y, x) + return tape.gradient(dy, x) + + for dtype in (np.float32, np.float64): + x = np.asarray( + [[-0.9, -0.7, -0.5, -0.3, -0.1], [0.1, 0.3, 0.5, 0.7, 0.9]], + dtype=dtype, + order="F", + ) + got = self.evaluate(f(constant_op.constant(x))) + want = _softsign_grad_grad(x.astype(np.float64)).astype(dtype) + self.assertAllClose(got, want) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/ops/nn_grad.py b/tensorflow/python/ops/nn_grad.py index a7b4a1c2b63e1f..1b7dcc0ac9314a 100644 --- a/tensorflow/python/ops/nn_grad.py +++ b/tensorflow/python/ops/nn_grad.py @@ -492,6 +492,20 @@ def _SoftsignGrad(op: ops.Operation, grad): return gen_nn_ops.softsign_grad(grad, op.inputs[0]) +@ops.RegisterGradient("SoftsignGrad") +def _SoftsignGradGrad(op: ops.Operation, grad): + x = op.inputs[1] + denominator = 1.0 + math_ops.abs(x) + return ( + gen_nn_ops.softsign_grad(grad, x), + -2.0 + * grad + * op.inputs[0] + * math_ops.sign(x) + / (denominator * denominator * denominator), + ) + + @ops.RegisterGradient("ReluGrad") def _ReluGradGrad(op: ops.Operation, grad): x = op.inputs[1] diff --git a/third_party/nccl/nccl_wheel.patch b/third_party/nccl/nccl_wheel.patch index f87b59ab0da5c2..30caf94090d360 100644 --- a/third_party/nccl/nccl_wheel.patch +++ b/third_party/nccl/nccl_wheel.patch @@ -1,48 +1,54 @@ diff --git a/include/nccl_device/gin/proxy/gin_proxy.h b/include/nccl_device/gin/proxy/gin_proxy.h +index 8a6d7dd2..78bd809e 100644 --- a/include/nccl_device/gin/proxy/gin_proxy.h +++ b/include/nccl_device/gin/proxy/gin_proxy.h -@@ -55,7 +55,7 @@ NCCL_DEVICE_INLINE void postGfd(Coop coop, ncclGinProxyGpuCtx_t* proxyCtx, ncclG --// 4x16 byte store with the write-through cache hint -+// 16x4 byte store with the write-through cache hint - #pragma unroll -- for (uint8_t i = 0; i < 4; i++) { -- __stwt((uint4*)&q[idx] + i, ((uint4*)gfd)[i]); +@@ -85,8 +85,8 @@ NCCL_DEVICE_INLINE void postGfd(Coop coop, ncclGinProxyGpuCtx_t* proxyCtx, ncclG + // ncclGinProxyGfd_t is declared __attribute__((packed, aligned(16))) in + // gin_proxy_device_host_common.h; static_asserts there enforce the contract. + NVCC_PRAGMA_UNROLL_AUTO +- for (uint8_t i = 0; i < sizeof(ncclGinProxyGfd_t) / sizeof(uint4); i++) { +- __stwt((uint4*)&q[gfdIdx] + i, ((uint4*)gfd)[i]); + for (uint8_t i = 0; i < 16; i++) { -+ __stwt((__half2*)&q[idx] + i, ((__half2*)gfd)[i]); ++ __stwt((__half2*)&q[gfdIdx] + i, ((__half2*)gfd)[i]); } - } - } + if (isGet) { + // Atomic max with rolling logic. diff --git a/include/nccl_device/utility.h b/include/nccl_device/utility.h +index 6e09fc10..dbc1cc28 100644 --- a/include/nccl_device/utility.h +++ b/include/nccl_device/utility.h @@ -17,7 +17,7 @@ - #define NCCL_CHECK_CUDACC 0 - #endif - #else -- #if __CUDACC__ -+ #ifdef __CUDACC__ - #define NCCL_CHECK_CUDACC 1 - #else - #define NCCL_CHECK_CUDACC 0 - #endif -diff --git a/include/nccl_device/gin/gdaki/doca_gpunetio/common/doca_gpunetio_verbs_def.h b/include/nccl_device/gin/gdaki/doca_gpunetio/common/doca_gpunetio_verbs_def.h ---- a/include/nccl_device/gin/gdaki/doca_gpunetio/common/doca_gpunetio_verbs_def.h -+++ b/include/nccl_device/gin/gdaki/doca_gpunetio/common/doca_gpunetio_verbs_def.h -@@ -49,7 +49,7 @@ extern "C" { - /** - * Macro to temporarily cast a variable to volatile. - */ --#define DOCA_GPUNETIO_VOLATILE(x) (*(volatile typeof(x) *)&(x)) -+#define DOCA_GPUNETIO_VOLATILE(x) (*(volatile __typeof__(x) *)&(x)) - - /** - * Default warp size value of 32 threads -@@ -123,7 +123,7 @@ extern "C" { - (1ULL << DOCA_GPUNETIO_VERBS_MAX_TRANSFER_SIZE_SHIFT) // 1GiB + #define NCCL_CHECK_CUDACC 0 + #endif + #else +-#if __CUDACC__ ++#ifdef __CUDACC__ + #define NCCL_CHECK_CUDACC 1 + #else + #define NCCL_CHECK_CUDACC 0 +diff --git a/include/nccl_device/gin/gpi/gin_gpi_device_host_common.h b/include/nccl_device/gin/gpi/gin_gpi_device_host_common.h +index 4a2cf812..e38f7fc2 100644 +--- a/include/nccl_device/gin/gpi/gin_gpi_device_host_common.h ++++ b/include/nccl_device/gin/gpi/gin_gpi_device_host_common.h +@@ -170,7 +170,7 @@ enum gpi_post_mode { + }; - #ifndef ACCESS_ONCE --#define ACCESS_ONCE(x) (*(volatile typeof(x) *)&(x)) -+#define ACCESS_ONCE(x) (*(volatile __typeof__(x) *)&(x)) + #ifndef GPI_ACCESS_ONCE +-#define GPI_ACCESS_ONCE(x) (*(volatile typeof(x)*)&(x)) ++#define GPI_ACCESS_ONCE(x) (*(volatile __typeof__(x)*)&(x)) #endif - #ifndef READ_ONCE \ No newline at end of file + #ifndef GPI_READ_ONCE +diff --git a/include/nccl_device/gin/gpi/gin_gpi.h b/include/nccl_device/gin/gpi/gin_gpi.h +index c004b4cb..7357c4c0 100644 +--- a/include/nccl_device/gin/gpi/gin_gpi.h ++++ b/include/nccl_device/gin/gpi/gin_gpi.h +@@ -212,7 +212,7 @@ __device__ static inline void gpi_gpu_channel_post_gfd_thread(gpi_gpu_channel_t* + // Manual PTX for MMIO 128-bit store (b128 needs CUDA 12.3+ / PTX 8.3) + uint64_t val_lo = segment[0].raw; + uint64_t val_hi = segment[1].raw; +-#if CUDART_VERSION >= 12030 ++#if CUDART_VERSION >= 12030 && !defined(__clang__) + asm volatile(R"YYY( + .reg .b128 _v%=; + mov.b128 _v%=, {%1, %2}; \ No newline at end of file diff --git a/third_party/nccl/workspace.bzl b/third_party/nccl/workspace.bzl index 38acb8b1693ce9..fb54ca1043572e 100644 --- a/third_party/nccl/workspace.bzl +++ b/third_party/nccl/workspace.bzl @@ -22,7 +22,7 @@ def repo(): name = "nccl_archive", build_file = "@xla//third_party/nccl:archive.BUILD", patch_file = ["@xla//third_party/nccl:archive.patch"], - sha256 = "e67239212c395bfdb398a7519491840d06fdf6b599c299f97c7ed0109777bba1", - strip_prefix = "nccl-2.29.7-1", - urls = tf_mirror_urls("https://github.com/NVIDIA/nccl/archive/refs/tags/v2.29.7-1.tar.gz"), + sha256 = "292a7f7a27b6754acaf46b5506a60758ca7b18cc1dfbd3d1d4e1e229d0863b4e", + strip_prefix = "nccl-2.30.7-1", + urls = tf_mirror_urls("https://github.com/NVIDIA/nccl/archive/refs/tags/v2.30.7-1.tar.gz"), ) diff --git a/third_party/xla/.github/workflows/clang_format.yml b/third_party/xla/.github/workflows/clang_format.yml index 477e70220a3bba..359d5b78d2a2d7 100644 --- a/third_party/xla/.github/workflows/clang_format.yml +++ b/third_party/xla/.github/workflows/clang_format.yml @@ -41,7 +41,7 @@ jobs: - name: "Fetch HEAD of main branch" run: git fetch origin main --depth=1 - name: Install uv dependencies - uses: astral-sh/setup-uv@ae62891fec2bb8e7d6c99fc78c9fec3a63790f8d # v10.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 - name: "Run clang-format check" env: TARGET_REF: "origin/${{ github.base_ref || 'main' }}" diff --git a/third_party/xla/.github/workflows/scorecards-analysis.yml b/third_party/xla/.github/workflows/scorecards-analysis.yml index 81077b90ea3b6f..f08174a83b51d0 100644 --- a/third_party/xla/.github/workflows/scorecards-analysis.yml +++ b/third_party/xla/.github/workflows/scorecards-analysis.yml @@ -67,6 +67,6 @@ jobs: # Upload the results to GitHub's code scanning dashboard (optional). # Commenting out will disable upload of results to your repo's Code Scanning dashboard - name: "Upload to code-scanning" - uses: github/codeql-action/upload-sarif@5595ccaf912efad79be6eef63a5619ff05969be3 # v4.37.6 + uses: github/codeql-action/upload-sarif@ff2f1c621b7f889edc0d3c761ac2e6a3f8cdb0dd # v4.37.7 with: sarif_file: results.sarif diff --git a/third_party/xla/.github/workflows/zizmor.yml b/third_party/xla/.github/workflows/zizmor.yml index 9268cac0007cf7..3bcae86a2acdb9 100644 --- a/third_party/xla/.github/workflows/zizmor.yml +++ b/third_party/xla/.github/workflows/zizmor.yml @@ -33,7 +33,7 @@ jobs: persist-credentials: false - name: Install uv dependencies - uses: astral-sh/setup-uv@ae62891fec2bb8e7d6c99fc78c9fec3a63790f8d # v10.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 - name: Run zizmor run: uvx zizmor --format=github .github xla diff --git a/third_party/xla/requirements_lock_3_11.txt b/third_party/xla/requirements_lock_3_11.txt index c130f9c07482a7..761e0ace834bd6 100644 --- a/third_party/xla/requirements_lock_3_11.txt +++ b/third_party/xla/requirements_lock_3_11.txt @@ -73,43 +73,44 @@ numpy==2.4.6 \ --hash=sha256:f74a575920ab21fe304421a3fc28793d82e299cae9eccb37084e9fc7f3617c20 lit==17.0.6 \ --hash=sha256:dfa9af9b55fc4509a56be7bf2346f079d7f4a242d583b9f2e0b078fd0abae31b -ml-dtypes==0.5.4 \ - --hash=sha256:0d2ffd05a2575b1519dc928c0b93c06339eb67173ff53acb00724502cda231cf \ - --hash=sha256:11942cbf2cf92157db91e5022633c0d9474d4dfd813a909383bd23ce828a4b7d \ - --hash=sha256:14a4fd3228af936461db66faccef6e4f41c1d82fcc30e9f8d58a08916b1d811f \ - --hash=sha256:19b9a53598f21e453ea2fbda8aa783c20faff8e1eeb0d7ab899309a0053f1483 \ - --hash=sha256:2314892cdc3fcf05e373d76d72aaa15fda9fb98625effa73c1d646f331fcecb7 \ - --hash=sha256:2b857d3af6ac0d39db1de7c706e69c7f9791627209c3d6dedbfca8c7e5faec22 \ - --hash=sha256:304ad47faa395415b9ccbcc06a0350800bc50eda70f0e45326796e27c62f18b6 \ - --hash=sha256:35f29491a3e478407f7047b8a4834e4640a77d2737e0b294d049746507af5175 \ - --hash=sha256:388d399a2152dd79a3f0456a952284a99ee5c93d3e2f8dfe25977511e0515270 \ - --hash=sha256:3bbbe120b915090d9dd1375e4684dd17a20a2491ef25d640a908281da85e73f1 \ - --hash=sha256:3d277bf3637f2a62176f4575512e9ff9ef51d00e39626d9fe4a161992f355af2 \ - --hash=sha256:4381fe2f2452a2d7589689693d3162e876b3ddb0a832cde7a414f8e1adf7eab1 \ - --hash=sha256:4ff7f3e7ca2972e7de850e7b8fcbb355304271e2933dd90814c1cb847414d6e2 \ - --hash=sha256:531eff30e4d368cb6255bc2328d070e35836aa4f282a0fb5f3a0cd7260257298 \ - --hash=sha256:533ce891ba774eabf607172254f2e7260ba5f57bdd64030c9a4fcfbd99815d0d \ - --hash=sha256:557a31a390b7e9439056644cb80ed0735a6e3e3bb09d67fd5687e4b04238d1de \ - --hash=sha256:5a0f68ca8fd8d16583dfa7793973feb86f2fbb56ce3966daf9c9f748f52a2049 \ - --hash=sha256:6a0df4223b514d799b8a1629c65ddc351b3efa833ccf7f8ea0cf654a61d1e35d \ - --hash=sha256:6c7ecb74c4bd71db68a6bea1edf8da8c34f3d9fe218f038814fd1d310ac76c90 \ - --hash=sha256:7c23c54a00ae43edf48d44066a7ec31e05fdc2eee0be2b8b50dd1903a1db94bb \ - --hash=sha256:805cef3a38f4eafae3a5bf9ebdcdb741d0bcfd9e1bd90eb54abd24f928cd2465 \ - --hash=sha256:88c982aac7cb1cbe8cbb4e7f253072b1df872701fcaf48d84ffbb433b6568f24 \ - --hash=sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453 \ - --hash=sha256:8c6a2dcebd6f3903e05d51960a8058d6e131fe69f952a5397e5dbabc841b6d56 \ - --hash=sha256:8c760d85a2f82e2bed75867079188c9d18dae2ee77c25a54d60e9cc79be1bc48 \ - --hash=sha256:9ad459e99793fa6e13bd5b7e6792c8f9190b4e5a1b45c63aba14a4d0a7f1d5ff \ - --hash=sha256:9bad06436568442575beb2d03389aa7456c690a5b05892c471215bfd8cf39460 \ - --hash=sha256:a174837a64f5b16cab6f368171a1a03a27936b31699d167684073ff1c4237dac \ - --hash=sha256:a7f7c643e8b1320fd958bf098aa7ecf70623a42ec5154e3be3be673f4c34d900 \ - --hash=sha256:a9b61c19040397970d18d7737375cffd83b1f36a11dd4ad19f83a016f736c3ef \ - --hash=sha256:b4b801ebe0b477be666696bda493a9be8356f1f0057a57f1e35cd26928823e5a \ - --hash=sha256:b95e97e470fe60ed493fd9ae3911d8da4ebac16bd21f87ffa2b7c588bf22ea2c \ - --hash=sha256:bc11d7e8c44a65115d05e2ab9989d1e045125d7be8e05a071a48bc76eb6d6040 \ - --hash=sha256:bfc534409c5d4b0bf945af29e5d0ab075eae9eecbb549ff8a29280db822f34f9 \ - --hash=sha256:c1a953995cccb9e25a4ae19e34316671e4e2edaebe4cf538229b1fc7109087b7 \ - --hash=sha256:cb73dccfc991691c444acc8c0012bee8f2470da826a92e3a20bb333b1a7894e6 \ - --hash=sha256:ce756d3a10d0c4067172804c9cc276ba9cc0ff47af9078ad439b075d1abdc29b \ - --hash=sha256:d81fdb088defa30eb37bf390bb7dde35d3a83ec112ac8e33d75ab28cc29dd8b0 \ - --hash=sha256:f21c9219ef48ca5ee78402d5cc831bd58ea27ce89beda894428bc67a52da5328 +ml-dtypes==0.6.0 \ + --hash=sha256:008382aeab529df5d3f00501ad9a7dcd64494d4b5b1971fc4c79019e6c1f5010 \ + --hash=sha256:03ce583adfce34ad33aa9e1fc7a8344dcf90ea776cc4ef0e5a48d4eae84e5d20 \ + --hash=sha256:084dfe51a7ad58b171f05115f8226ed4233a454a1611371947e806e76f0c638d \ + --hash=sha256:26b1f1fa4f0435a2946859823f6e2bf06796f1e9f10f5a05b08a5e3c8f46ff69 \ + --hash=sha256:28d676428b104bb9717b0928bc5c5129f2d6b51b6727587cc4289e7bf8713cb5 \ + --hash=sha256:2a3e9d53925597fbffafd2a37048dadeddd0bdaba58058f6ae0869ed709a184d \ + --hash=sha256:3035518e3e19add1a4cac9236ab22888b208a4074912514313ccb2d6d242cde8 \ + --hash=sha256:317be9967fb84b0ce4e80e6b1bf71213d21971621cf6f1e501a63602a95297bf \ + --hash=sha256:31f1ce979d31a357e95aa81812f20412c8c954fa43c44ee3ead1e1c8a78575ef \ + --hash=sha256:37da32aa97749251025666d62372775019594577b9c9e9cfda83bed48d778fdb \ + --hash=sha256:3b4a480aa8fd54a1805b8ac10f3f91763926a74f73c0c364c10f9231854f4170 \ + --hash=sha256:3be9911d953f97cddded4b9961d7b650473b7e55806d20f6176f8356dfe7b38e \ + --hash=sha256:3e169214e0d80ff1c038e1b3017e33c23e43bdf948d42d31de8283111c7e2fa3 \ + --hash=sha256:488c99ab181a2f59d9ec3b12c5fa11ec904e92be2c4ba18cded54dd7501208fe \ + --hash=sha256:5359c588cc62de6f78d7430f06b65853d884955494d86d6ad90b6dd64a3f3a08 \ + --hash=sha256:573b11f3c327e17ef3826d266e676cf1149a1f3016f822a05f2306c55d8246bf \ + --hash=sha256:57ed0d6b4ac5e7868361303a9c57fbcf63b768236ee14456f585dfcf260d0292 \ + --hash=sha256:5a519c9e95a216fbcb8e759793ef7fb40793fc803ed839142d6dc5be9be5bc89 \ + --hash=sha256:5e60251d32ced5598972e4d5e06a2f044341f9291402551a3f6f0ec44f9299b0 \ + --hash=sha256:6c8e39b53e90afda8ce52859c93de4dba3e02b76d85dcf091cc469f9184c6dae \ + --hash=sha256:6eaed129a4afe90694b8685e2f9b6294849f5eda4af9a15be83a4326eeebd775 \ + --hash=sha256:6ec0d244a5bba12239025389ad88bbfb45f9f10e25ab4f678e9a4768ebd47532 \ + --hash=sha256:7728c0420ec1c338564fc8b01015ff2d58567e70f17fedce5a0a7c0308c0d5b9 \ + --hash=sha256:84fa136b8602c8c39e3b6cb24918960cd6f36cade7a70376f56770729cd56510 \ + --hash=sha256:8f490c003369ce60e514a0c3b12374f05274c101fee1bead6740ec8a564032b0 \ + --hash=sha256:9c6ad60af4102789a5c09824004beade2f7f28cd1cd581ee5c170d9dc2fbb00e \ + --hash=sha256:b1b503864fada3f74fabf8d9fee7b4c1cbe956301e6fdece975d5f77c2fce958 \ + --hash=sha256:b76fa1d3f92967d58289ac47ab7458ede66e6f3527fff3e59142aee57d9307cd \ + --hash=sha256:bad8d1dd5bed060a29332b99d63d0e5c2969081e1c6ea54adfbccfdfa783be44 \ + --hash=sha256:ce7563e0b1a4482cbc1b4a6272145e54e4489e54fe7428f94908c3d87103abfa \ + --hash=sha256:d4f1b9329a251e4affe3bb58f4d3e2db22a714396fd7ffb40d0b5db423c24d17 \ + --hash=sha256:d574c2b28921dc72e869df248f1a278f6eee176a1f237c8642e1a71eb15f3977 \ + --hash=sha256:de9d14748dbf3968951436ef514a29c9d1fe438aa680d110134ee2f7a9f9df18 \ + --hash=sha256:e25bb3b0ad1217b60626e4ed45b10ca170c41d99fbe44a12bebc1e07ec4aad55 \ + --hash=sha256:e2d6149f3a57f405bcad5fb41e03218b8373936253f23e1ca84c0108abbc3392 \ + --hash=sha256:e74266ca8e97874a937b7646378c178025650a236584f7474d10d8086a6edea3 \ + --hash=sha256:f4adb4af61516510d786cf8c01851a66f6d3ddfa79e1144deaa5b40d8507231e \ + --hash=sha256:f4f59f83c82ab480e924b988e7b1b4eb4de836dfcf5390c6f59148d1a00e1d02 \ + --hash=sha256:f6cb525101b6b903779188c1e9e9490c343b455ab822883e02cf01e5547338d2 \ + --hash=sha256:fb87f46b4f7ad7b5d3ad8f4b452b024bd4229d44c8ff934798c1fe656210387a diff --git a/third_party/xla/requirements_lock_3_12.txt b/third_party/xla/requirements_lock_3_12.txt index e6fb1e433c71f7..097af0257f1d12 100644 --- a/third_party/xla/requirements_lock_3_12.txt +++ b/third_party/xla/requirements_lock_3_12.txt @@ -45,43 +45,44 @@ numpy==2.5.1 \ --hash=sha256:f7feb014281029e628ba2d5a007407443b06e418b6fe451d1e2adcbc8eba0107 lit==17.0.6 \ --hash=sha256:dfa9af9b55fc4509a56be7bf2346f079d7f4a242d583b9f2e0b078fd0abae31b -ml-dtypes==0.5.4 \ - --hash=sha256:0d2ffd05a2575b1519dc928c0b93c06339eb67173ff53acb00724502cda231cf \ - --hash=sha256:11942cbf2cf92157db91e5022633c0d9474d4dfd813a909383bd23ce828a4b7d \ - --hash=sha256:14a4fd3228af936461db66faccef6e4f41c1d82fcc30e9f8d58a08916b1d811f \ - --hash=sha256:19b9a53598f21e453ea2fbda8aa783c20faff8e1eeb0d7ab899309a0053f1483 \ - --hash=sha256:2314892cdc3fcf05e373d76d72aaa15fda9fb98625effa73c1d646f331fcecb7 \ - --hash=sha256:2b857d3af6ac0d39db1de7c706e69c7f9791627209c3d6dedbfca8c7e5faec22 \ - --hash=sha256:304ad47faa395415b9ccbcc06a0350800bc50eda70f0e45326796e27c62f18b6 \ - --hash=sha256:35f29491a3e478407f7047b8a4834e4640a77d2737e0b294d049746507af5175 \ - --hash=sha256:388d399a2152dd79a3f0456a952284a99ee5c93d3e2f8dfe25977511e0515270 \ - --hash=sha256:3bbbe120b915090d9dd1375e4684dd17a20a2491ef25d640a908281da85e73f1 \ - --hash=sha256:3d277bf3637f2a62176f4575512e9ff9ef51d00e39626d9fe4a161992f355af2 \ - --hash=sha256:4381fe2f2452a2d7589689693d3162e876b3ddb0a832cde7a414f8e1adf7eab1 \ - --hash=sha256:4ff7f3e7ca2972e7de850e7b8fcbb355304271e2933dd90814c1cb847414d6e2 \ - --hash=sha256:531eff30e4d368cb6255bc2328d070e35836aa4f282a0fb5f3a0cd7260257298 \ - --hash=sha256:533ce891ba774eabf607172254f2e7260ba5f57bdd64030c9a4fcfbd99815d0d \ - --hash=sha256:557a31a390b7e9439056644cb80ed0735a6e3e3bb09d67fd5687e4b04238d1de \ - --hash=sha256:5a0f68ca8fd8d16583dfa7793973feb86f2fbb56ce3966daf9c9f748f52a2049 \ - --hash=sha256:6a0df4223b514d799b8a1629c65ddc351b3efa833ccf7f8ea0cf654a61d1e35d \ - --hash=sha256:6c7ecb74c4bd71db68a6bea1edf8da8c34f3d9fe218f038814fd1d310ac76c90 \ - --hash=sha256:7c23c54a00ae43edf48d44066a7ec31e05fdc2eee0be2b8b50dd1903a1db94bb \ - --hash=sha256:805cef3a38f4eafae3a5bf9ebdcdb741d0bcfd9e1bd90eb54abd24f928cd2465 \ - --hash=sha256:88c982aac7cb1cbe8cbb4e7f253072b1df872701fcaf48d84ffbb433b6568f24 \ - --hash=sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453 \ - --hash=sha256:8c6a2dcebd6f3903e05d51960a8058d6e131fe69f952a5397e5dbabc841b6d56 \ - --hash=sha256:8c760d85a2f82e2bed75867079188c9d18dae2ee77c25a54d60e9cc79be1bc48 \ - --hash=sha256:9ad459e99793fa6e13bd5b7e6792c8f9190b4e5a1b45c63aba14a4d0a7f1d5ff \ - --hash=sha256:9bad06436568442575beb2d03389aa7456c690a5b05892c471215bfd8cf39460 \ - --hash=sha256:a174837a64f5b16cab6f368171a1a03a27936b31699d167684073ff1c4237dac \ - --hash=sha256:a7f7c643e8b1320fd958bf098aa7ecf70623a42ec5154e3be3be673f4c34d900 \ - --hash=sha256:a9b61c19040397970d18d7737375cffd83b1f36a11dd4ad19f83a016f736c3ef \ - --hash=sha256:b4b801ebe0b477be666696bda493a9be8356f1f0057a57f1e35cd26928823e5a \ - --hash=sha256:b95e97e470fe60ed493fd9ae3911d8da4ebac16bd21f87ffa2b7c588bf22ea2c \ - --hash=sha256:bc11d7e8c44a65115d05e2ab9989d1e045125d7be8e05a071a48bc76eb6d6040 \ - --hash=sha256:bfc534409c5d4b0bf945af29e5d0ab075eae9eecbb549ff8a29280db822f34f9 \ - --hash=sha256:c1a953995cccb9e25a4ae19e34316671e4e2edaebe4cf538229b1fc7109087b7 \ - --hash=sha256:cb73dccfc991691c444acc8c0012bee8f2470da826a92e3a20bb333b1a7894e6 \ - --hash=sha256:ce756d3a10d0c4067172804c9cc276ba9cc0ff47af9078ad439b075d1abdc29b \ - --hash=sha256:d81fdb088defa30eb37bf390bb7dde35d3a83ec112ac8e33d75ab28cc29dd8b0 \ - --hash=sha256:f21c9219ef48ca5ee78402d5cc831bd58ea27ce89beda894428bc67a52da5328 +ml-dtypes==0.6.0 \ + --hash=sha256:008382aeab529df5d3f00501ad9a7dcd64494d4b5b1971fc4c79019e6c1f5010 \ + --hash=sha256:03ce583adfce34ad33aa9e1fc7a8344dcf90ea776cc4ef0e5a48d4eae84e5d20 \ + --hash=sha256:084dfe51a7ad58b171f05115f8226ed4233a454a1611371947e806e76f0c638d \ + --hash=sha256:26b1f1fa4f0435a2946859823f6e2bf06796f1e9f10f5a05b08a5e3c8f46ff69 \ + --hash=sha256:28d676428b104bb9717b0928bc5c5129f2d6b51b6727587cc4289e7bf8713cb5 \ + --hash=sha256:2a3e9d53925597fbffafd2a37048dadeddd0bdaba58058f6ae0869ed709a184d \ + --hash=sha256:3035518e3e19add1a4cac9236ab22888b208a4074912514313ccb2d6d242cde8 \ + --hash=sha256:317be9967fb84b0ce4e80e6b1bf71213d21971621cf6f1e501a63602a95297bf \ + --hash=sha256:31f1ce979d31a357e95aa81812f20412c8c954fa43c44ee3ead1e1c8a78575ef \ + --hash=sha256:37da32aa97749251025666d62372775019594577b9c9e9cfda83bed48d778fdb \ + --hash=sha256:3b4a480aa8fd54a1805b8ac10f3f91763926a74f73c0c364c10f9231854f4170 \ + --hash=sha256:3be9911d953f97cddded4b9961d7b650473b7e55806d20f6176f8356dfe7b38e \ + --hash=sha256:3e169214e0d80ff1c038e1b3017e33c23e43bdf948d42d31de8283111c7e2fa3 \ + --hash=sha256:488c99ab181a2f59d9ec3b12c5fa11ec904e92be2c4ba18cded54dd7501208fe \ + --hash=sha256:5359c588cc62de6f78d7430f06b65853d884955494d86d6ad90b6dd64a3f3a08 \ + --hash=sha256:573b11f3c327e17ef3826d266e676cf1149a1f3016f822a05f2306c55d8246bf \ + --hash=sha256:57ed0d6b4ac5e7868361303a9c57fbcf63b768236ee14456f585dfcf260d0292 \ + --hash=sha256:5a519c9e95a216fbcb8e759793ef7fb40793fc803ed839142d6dc5be9be5bc89 \ + --hash=sha256:5e60251d32ced5598972e4d5e06a2f044341f9291402551a3f6f0ec44f9299b0 \ + --hash=sha256:6c8e39b53e90afda8ce52859c93de4dba3e02b76d85dcf091cc469f9184c6dae \ + --hash=sha256:6eaed129a4afe90694b8685e2f9b6294849f5eda4af9a15be83a4326eeebd775 \ + --hash=sha256:6ec0d244a5bba12239025389ad88bbfb45f9f10e25ab4f678e9a4768ebd47532 \ + --hash=sha256:7728c0420ec1c338564fc8b01015ff2d58567e70f17fedce5a0a7c0308c0d5b9 \ + --hash=sha256:84fa136b8602c8c39e3b6cb24918960cd6f36cade7a70376f56770729cd56510 \ + --hash=sha256:8f490c003369ce60e514a0c3b12374f05274c101fee1bead6740ec8a564032b0 \ + --hash=sha256:9c6ad60af4102789a5c09824004beade2f7f28cd1cd581ee5c170d9dc2fbb00e \ + --hash=sha256:b1b503864fada3f74fabf8d9fee7b4c1cbe956301e6fdece975d5f77c2fce958 \ + --hash=sha256:b76fa1d3f92967d58289ac47ab7458ede66e6f3527fff3e59142aee57d9307cd \ + --hash=sha256:bad8d1dd5bed060a29332b99d63d0e5c2969081e1c6ea54adfbccfdfa783be44 \ + --hash=sha256:ce7563e0b1a4482cbc1b4a6272145e54e4489e54fe7428f94908c3d87103abfa \ + --hash=sha256:d4f1b9329a251e4affe3bb58f4d3e2db22a714396fd7ffb40d0b5db423c24d17 \ + --hash=sha256:d574c2b28921dc72e869df248f1a278f6eee176a1f237c8642e1a71eb15f3977 \ + --hash=sha256:de9d14748dbf3968951436ef514a29c9d1fe438aa680d110134ee2f7a9f9df18 \ + --hash=sha256:e25bb3b0ad1217b60626e4ed45b10ca170c41d99fbe44a12bebc1e07ec4aad55 \ + --hash=sha256:e2d6149f3a57f405bcad5fb41e03218b8373936253f23e1ca84c0108abbc3392 \ + --hash=sha256:e74266ca8e97874a937b7646378c178025650a236584f7474d10d8086a6edea3 \ + --hash=sha256:f4adb4af61516510d786cf8c01851a66f6d3ddfa79e1144deaa5b40d8507231e \ + --hash=sha256:f4f59f83c82ab480e924b988e7b1b4eb4de836dfcf5390c6f59148d1a00e1d02 \ + --hash=sha256:f6cb525101b6b903779188c1e9e9490c343b455ab822883e02cf01e5547338d2 \ + --hash=sha256:fb87f46b4f7ad7b5d3ad8f4b452b024bd4229d44c8ff934798c1fe656210387a diff --git a/third_party/xla/tensorflow.bazelrc b/third_party/xla/tensorflow.bazelrc index 054ade6bdf69c5..0c18ed64c47782 100644 --- a/third_party/xla/tensorflow.bazelrc +++ b/third_party/xla/tensorflow.bazelrc @@ -189,7 +189,7 @@ common:mkl_aarch64 --config=mkl_aarch64_threadpool common:cuda_version --repo_env=HERMETIC_CUDA_VERSION="13.2.0" common:cuda_version --repo_env=HERMETIC_CUDNN_VERSION="9.12.0" common:cuda_version --repo_env=HERMETIC_NVSHMEM_VERSION="3.3.20" -common:cuda_version --repo_env=HERMETIC_NCCL_VERSION="2.29.7" +common:cuda_version --repo_env=HERMETIC_NCCL_VERSION="2.30.7" # CUDA: This config refers to building CUDA op kernels with nvcc. common:cuda --repo_env TF_NEED_CUDA=1 diff --git a/third_party/xla/third_party/nccl/archive.BUILD b/third_party/xla/third_party/nccl/archive.BUILD index 5c8c2d0fd09687..70bb40f1017a2c 100644 --- a/third_party/xla/third_party/nccl/archive.BUILD +++ b/third_party/xla/third_party/nccl/archive.BUILD @@ -38,7 +38,7 @@ exports_files(["LICENSE.txt"]) NCCL_MAJOR = 2 -NCCL_MINOR = 29 +NCCL_MINOR = 30 NCCL_PATCH = 7 @@ -232,7 +232,6 @@ cc_library( # from the virtual includes directory. "src/include/collectives.h", "src/nccl.h", - "src/ras/ras_internal.h", ], hdrs = ["src/nccl.h"], include_prefix = "third_party/nccl", diff --git a/third_party/xla/third_party/nccl/archive.patch b/third_party/xla/third_party/nccl/archive.patch index 4b56bb6ca368d6..5844ea08fb9cab 100644 --- a/third_party/xla/third_party/nccl/archive.patch +++ b/third_party/xla/third_party/nccl/archive.patch @@ -2,29 +2,30 @@ diff --git a/src/device/common.cu b/src/device/common.cu.cc similarity index 100% rename from src/device/common.cu rename to src/device/common.cu.cc -diff --git a/src/device/onerank.cu b/src/device/onerank.cu.cc -similarity index 100% -rename from src/device/onerank.cu -rename to src/device/onerank.cu.cc diff --git a/src/device/common.h b/src/device/common.h +index 1c1f5d6f..700c27a7 100644 --- a/src/device/common.h +++ b/src/device/common.h -@@ -24,7 +24,7 @@ - #endif - - typedef void(*ncclDevFuncPtr_t)(); +@@ -26,9 +26,9 @@ + typedef void (*ncclDevFuncPtr_t)(); + #if defined(NCCL_OS_WINDOWS) + /* MSVC C2133: extern array of unknown size needs a complete type; use pointer instead. */ +-extern __device__ ncclDevFuncPtr_t const* ncclDevFuncTable; ++extern __device__ ncclDevFuncPtr_t* ncclDevFuncTable; + #else -extern __device__ ncclDevFuncPtr_t const ncclDevFuncTable[]; +extern __device__ ncclDevFuncPtr_t ncclDevFuncTable[]; - + #endif + struct ncclShmemGroup { - ncclConnInfo *recvConns[NCCL_MAX_ARITY]; diff --git a/src/device/generate.py b/src/device/generate.py +index 88ddcef9..73b88d2b 100755 --- a/src/device/generate.py +++ b/src/device/generate.py -@@ -209,8 +209,8 @@ kernel_funcs = sorted(set(best_kernel(*fn) for fn in primary_funcs)) - +@@ -207,8 +207,8 @@ kernel_funcs = sorted(set(best_kernel(*fn) for fn in primary_funcs)) + ################################################################################ - + -# Generate /device_table.cu -with open(os.path.join(gensrc, "device_table.cu"), "w") as f: +# Generate /device_table.cu.cc @@ -32,16 +33,19 @@ diff --git a/src/device/generate.py b/src/device/generate.py out = f.write out('#include "common.h"\n') out("\n") -@@ -225,7 +225,7 @@ with open(os.path.join(gensrc, "device_table.cu"), "w") as f: - out("#endif\n") - out("\n") - -- out("__device__ ncclDevFuncPtr_t const ncclDevFuncTable[] = {\n"); -+ out("__device__ ncclDevFuncPtr_t ncclDevFuncTable[] = {\n"); +@@ -227,9 +227,9 @@ with open(os.path.join(gensrc, "device_table.cu"), "w") as f: + # an internal array + pointer alias. On Linux/GCC the array is named directly + # to avoid an extra device-memory indirection on every kernel dispatch. + out("#if defined(NCCL_OS_WINDOWS)\n") +- out("__device__ ncclDevFuncPtr_t const ncclDevFuncTableData[] = {\n") ++ out("__device__ ncclDevFuncPtr_t ncclDevFuncTableData[] = {\n") + out("#else\n") +- out("__device__ ncclDevFuncPtr_t const ncclDevFuncTable[] = {\n") ++ out("__device__ ncclDevFuncPtr_t ncclDevFuncTable[] = {\n") + out("#endif\n") index = 0 for fn in primary_funcs: - sym = paste("_", "ncclDevFunc", *fn) -@@ -283,8 +283,16 @@ with open(os.path.join(gensrc, "host_table.cc"), "w") as f: +@@ -292,8 +292,16 @@ with open(os.path.join(gensrc, "host_table.cc"), "w") as f: cudart, _ = required_cuda(*kfn) sym = paste("_", "ncclDevKernel", *kfn) if cudart != 0: out("#if CUDART_VERSION >= %d\n" % cudart) @@ -60,9 +64,9 @@ diff --git a/src/device/generate.py b/src/device/generate.py index += 1 out("nullptr};\n") out("\n") -@@ -298,15 +306,23 @@ with open(os.path.join(gensrc, "host_table.cc"), "w") as f: +@@ -307,15 +315,23 @@ with open(os.path.join(gensrc, "host_table.cc"), "w") as f: out("\n") - + # Maps primary id to kernel function pointer. - out("extern void* const ncclDevKernelForFunc[] = {\n") index = 0 @@ -87,16 +91,16 @@ diff --git a/src/device/generate.py b/src/device/generate.py index += 1 out("nullptr};\n") out("\n") -@@ -325,7 +341,7 @@ with open(os.path.join(gensrc, "host_table.cc"), "w") as f: +@@ -334,7 +350,7 @@ with open(os.path.join(gensrc, "host_table.cc"), "w") as f: # "coll" is reflected in the name: formally that no two funcs having different # coll's map to the same filename. def impl_filename(coll, redop, ty, algo, proto): - return "%s.cu" % paste("_", coll_camel_to_lower[coll], redop and redop.lower(), ty) + return "%s.cu.cc" % paste("_", coll_camel_to_lower[coll], redop and redop.lower(), ty) - + # Partition the functions and kernels to the .cu filenames. The partition is # a dictionary mapping filename to (coll, func-tuple list) -@@ -357,7 +373,7 @@ if os.environ.get("NCCL_USE_CMAKE", "0") != "1": +@@ -366,7 +382,7 @@ if os.environ.get("NCCL_USE_CMAKE", "0") != "1": with open(os.path.join(gensrc, "rules.mk"), "w") as f: out = f.write impl_names = sorted(name_to_funcs.keys()) @@ -105,66 +109,71 @@ diff --git a/src/device/generate.py b/src/device/generate.py out("LIB_OBJS_GEN = $(patsubst %,$(OBJDIR)/genobj/%.o,{names})\n" .format(names=" ".join(names))) out("\n") - +diff --git a/src/device/onerank.cu b/src/device/onerank.cu.cc +similarity index 100% +rename from src/device/onerank.cu +rename to src/device/onerank.cu.cc +diff --git a/src/device/symmetric/data_ops.cuh b/src/device/symmetric/data_ops.cuh +index 25b6a922..315b04d4 100644 +--- a/src/device/symmetric/data_ops.cuh ++++ b/src/device/symmetric/data_ops.cuh +@@ -41,7 +41,7 @@ static __device__ __forceinline__ T loadMem(GMemTag, T* p) { + for (int i = 0; i < sizeof(T) / 4; i++) u32[i] = __ldcs((uint32_t*)p + i); + break; + case 8: +- for (int i = 0; i < sizeof(T) / 8; i++) u64[i] = __ldcs((uint64_t*)p + i); ++ for (int i = 0; i < sizeof(T) / 8; i++) u64[i] = __ldcs((unsigned long long*)p + i); + break; + case 16: + for (int i = 0; i < sizeof(T) / 16; i++) u32v4[i] = __ldcs((uint4*)p + i); +@@ -89,7 +89,7 @@ static __device__ __forceinline__ void storeMem(GMemTag, T* p, T val) { + for (int i = 0; i < sizeof(T) / 4; i++) __stcs((uint32_t*)p + i, u32[i]); + break; + case 8: +- for (int i = 0; i < sizeof(T) / 8; i++) __stcs((uint64_t*)p + i, u64[i]); ++ for (int i = 0; i < sizeof(T) / 8; i++) __stcs((unsigned long long*)p + i, u64[i]); + break; + case 16: + for (int i = 0; i < sizeof(T) / 16; i++) __stcs((uint4*)p + i, u32v4[i]); diff --git a/src/include/nccl_common.h b/src/include/nccl_common.h +index 70b669a5..2667e055 100644 --- a/src/include/nccl_common.h +++ b/src/include/nccl_common.h -@@ -9,7 +9,9 @@ - #define NCCL_DEBUG_H_ - - // Workaround for libstdc++ trying to force public visibility of std:: symbols. We don't want to do that in libnccl.so. +@@ -11,7 +11,9 @@ + #ifdef NCCL_OS_LINUX + // Workaround for libstdc++ trying to force public visibility of std:: symbols. We don't want to do that in + // libnccl.so. +#if defined(__GLIBCXX__) #include +#endif #undef _GLIBCXX_VISIBILITY #define _GLIBCXX_VISIBILITY(V) - + #endif diff --git a/src/include/nccl_device/gin/proxy/gin_proxy.h b/src/include/nccl_device/gin/proxy/gin_proxy.h +index 8a6d7dd2..78bd809e 100644 --- a/src/include/nccl_device/gin/proxy/gin_proxy.h +++ b/src/include/nccl_device/gin/proxy/gin_proxy.h -@@ -55,7 +55,7 @@ NCCL_DEVICE_INLINE void postGfd(Coop coop, ncclGinProxyGpuCtx_t* proxyCtx, ncclG --// 4x16 byte store with the write-through cache hint -+// 16x4 byte store with the write-through cache hint - #pragma unroll -- for (uint8_t i = 0; i < 4; i++) { -- __stwt((uint4*)&q[idx] + i, ((uint4*)gfd)[i]); +@@ -85,8 +85,8 @@ NCCL_DEVICE_INLINE void postGfd(Coop coop, ncclGinProxyGpuCtx_t* proxyCtx, ncclG + // ncclGinProxyGfd_t is declared __attribute__((packed, aligned(16))) in + // gin_proxy_device_host_common.h; static_asserts there enforce the contract. + NVCC_PRAGMA_UNROLL_AUTO +- for (uint8_t i = 0; i < sizeof(ncclGinProxyGfd_t) / sizeof(uint4); i++) { +- __stwt((uint4*)&q[gfdIdx] + i, ((uint4*)gfd)[i]); + for (uint8_t i = 0; i < 16; i++) { -+ __stwt((__half2*)&q[idx] + i, ((__half2*)gfd)[i]); ++ __stwt((__half2*)&q[gfdIdx] + i, ((__half2*)gfd)[i]); } - } - } -diff --git a/src/include/nccl_device/utility.h b/src/include/nccl_device/utility.h ---- a/src/include/nccl_device/utility.h -+++ b/src/include/nccl_device/utility.h -@@ -17,7 +17,7 @@ - #define NCCL_CHECK_CUDACC 0 - #endif - #else -- #if __CUDACC__ -+ #ifdef __CUDACC__ - #define NCCL_CHECK_CUDACC 1 - #else - #define NCCL_CHECK_CUDACC 0 - #endif - -diff --git a/src/device/symmetric/data_ops.cuh b/src/device/symmetric/data_ops.cuh ---- a/src/device/symmetric/data_ops.cuh -+++ b/src/device/symmetric/data_ops.cuh -@@ -26,7 +26,7 @@ - case 1: for (int i=0; i < sizeof(T)/1; i++) u8[i] = __ldcs((uint8_t*)p + i); break; - case 2: for (int i=0; i < sizeof(T)/2; i++) u16[i] = __ldcs((uint16_t*)p + i); break; - case 4: for (int i=0; i < sizeof(T)/4; i++) u32[i] = __ldcs((uint32_t*)p + i); break; -- case 8: for (int i=0; i < sizeof(T)/8; i++) u64[i] = __ldcs((uint64_t*)p + i); break; -+ case 8: for (int i=0; i < sizeof(T)/8; i++) u64[i] = __ldcs((unsigned long long*)p + i); break; - case 16: for (int i=0; i < sizeof(T)/16; i++) u32v4[i] = __ldcs((uint4*)p + i); break; - default: __builtin_unreachable(); - } -@@ -62,7 +62,7 @@ - case 1: for (int i=0; i < sizeof(T)/1; i++) __stcs((uint8_t*)p + i, u8[i]); break; - case 2: for (int i=0; i < sizeof(T)/2; i++) __stcs((uint16_t*)p + i, u16[i]); break; - case 4: for (int i=0; i < sizeof(T)/4; i++) __stcs((uint32_t*)p + i, u32[i]); break; -- case 8: for (int i=0; i < sizeof(T)/8; i++) __stcs((uint64_t*)p + i, u64[i]); break; -+ case 8: for (int i=0; i < sizeof(T)/8; i++) __stcs((unsigned long long*)p + i, u64[i]); break; - case 16: for (int i=0; i < sizeof(T)/16; i++) __stcs((uint4*)p + i, u32v4[i]); break; - default: __builtin_unreachable(); - } + if (isGet) { + // Atomic max with rolling logic. +diff --git a/src/include/nccl_device/gin/gpi/gin_gpi.h b/src/include/nccl_device/gin/gpi/gin_gpi.h +index c004b4cb..7357c4c0 100644 +--- a/src/include/nccl_device/gin/gpi/gin_gpi.h ++++ b/src/include/nccl_device/gin/gpi/gin_gpi.h +@@ -212,7 +212,7 @@ __device__ static inline void gpi_gpu_channel_post_gfd_thread(gpi_gpu_channel_t* + // Manual PTX for MMIO 128-bit store (b128 needs CUDA 12.3+ / PTX 8.3) + uint64_t val_lo = segment[0].raw; + uint64_t val_hi = segment[1].raw; +-#if CUDART_VERSION >= 12030 ++#if CUDART_VERSION >= 12030 && !defined(__clang__) + asm volatile(R"YYY( + .reg .b128 _v%=; + mov.b128 _v%=, {%1, %2}; \ No newline at end of file diff --git a/third_party/xla/third_party/nccl/workspace.bzl b/third_party/xla/third_party/nccl/workspace.bzl index dd40c3264966d7..ba05442b846048 100644 --- a/third_party/xla/third_party/nccl/workspace.bzl +++ b/third_party/xla/third_party/nccl/workspace.bzl @@ -22,7 +22,7 @@ def repo(): name = "nccl_archive", build_file = "//third_party/nccl:archive.BUILD", patch_file = ["//third_party/nccl:archive.patch"], - sha256 = "e67239212c395bfdb398a7519491840d06fdf6b599c299f97c7ed0109777bba1", - strip_prefix = "nccl-2.29.7-1", - urls = tf_mirror_urls("https://github.com/NVIDIA/nccl/archive/refs/tags/v2.29.7-1.tar.gz"), + sha256 = "292a7f7a27b6754acaf46b5506a60758ca7b18cc1dfbd3d1d4e1e229d0863b4e", + strip_prefix = "nccl-2.30.7-1", + urls = tf_mirror_urls("https://github.com/NVIDIA/nccl/archive/refs/tags/v2.30.7-1.tar.gz"), ) diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch index fc07bbc2c695a5..dc2c4735acf69c 100644 --- a/third_party/xla/third_party/stablehlo/temporary.patch +++ b/third_party/xla/third_party/stablehlo/temporary.patch @@ -1,3 +1,41 @@ +diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md +--- stablehlo/docs/spec.md ++++ stablehlo/docs/spec.md +@@ -2493,7 +2493,9 @@ + * `num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1`. + * (C26) `rank(result) = N`. + * If the operation uses non-quantized tensors: +- * (C27) `element_type(lhs) = element_type(rhs) = element_type(result)`. ++ * (C27) `element_type(lhs) = element_type(rhs) = element_type(result)` or ++ (`is_fp8(element_type(lhs))` and `is_fp8(element_type(rhs))` and ++ `element_type(result) = element_type(lhs)`). + * If the operation uses quantized tensors: + * (C28) `is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)`. + * (C29) If `is_per_axis_quantized(rhs)`, +@@ -3163,7 +3165,9 @@ + * `num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1`. + * (C26) `rank(result) = N`. + * If the operation uses non-quantized tensors: +- * (C27) `element_type(lhs) = element_type(rhs) = element_type(result)`. ++ * (C27) `element_type(lhs) = element_type(rhs) = element_type(result)` or ++ (`is_fp8(element_type(lhs))` and `is_fp8(element_type(rhs))` and ++ `element_type(result) = element_type(lhs)`). + * If the operation uses quantized tensors: + * (C28) `is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)`. + * (C29) If `is_per_axis_quantized(rhs)`, +@@ -7654,6 +7658,12 @@ + * `is_quantized(x: Value | Placeholder | Type) -> Value` is a shortcut for + `is_quantized_tensor_element_type(x)`. + ++* `is_fp8(x: Value | Placeholder | Type) -> Value` returns `true` if `x` is one ++of `Float8E4M3Type`, `Float8E4M3FNType`, `Float8E4M3B11FNUZType`, ++`Float8E4M3FNUZType`, `Float8E5M2Type`, `Float8E5M2FNUZType`, `Float8E3M4Type`, ++or `Float8E8M0FNUType`. If `x` is a value or placeholder, this function is a ++shortcut for `is_fp8(type(x))`. ++ + * `is_type_name(x: Value | Placeholder | Type) -> Value`. Available for all + types. For example, `is_float(x)` returns `true` if `x` is a `FloatType`. + If `x` is a value or placeholder, this function is a shortcut for diff --ruN a/stablehlo/stablehlo/dialect/Serialization.cpp b/stablehlo/stablehlo/dialect/Serialization.cpp --- stablehlo/stablehlo/dialect/Serialization.cpp +++ stablehlo/stablehlo/dialect/Serialization.cpp @@ -67,6 +105,144 @@ diff --ruN a/stablehlo/stablehlo/dialect/Serialization.cpp b/stablehlo/stablehlo return nullptr; } +diff --ruN a/stablehlo/stablehlo/dialect/TypeInference.cpp b/stablehlo/stablehlo/dialect/TypeInference.cpp +--- stablehlo/stablehlo/dialect/TypeInference.cpp ++++ stablehlo/stablehlo/dialect/TypeInference.cpp +@@ -86,6 +86,13 @@ + // Utils for quantization specific verifications + //===----------------------------------------------------------------------===// + ++bool isFp8Type(mlir::Type type) { ++ return llvm::isa(type); ++} ++ + template + bool allQuantized(ArrayRef typeRange) { + return llvm::all_of( +@@ -2221,11 +2228,14 @@ + // convolution_c27 + if (!anyQuantized({rankedLhsType, rankedRhsType}) && + !isCompatibleForHloTypeInference(rankedLhsType.getElementType(), +- rankedRhsType.getElementType())) ++ rankedRhsType.getElementType()) && ++ !(isFp8Type(rankedLhsType.getElementType()) && ++ isFp8Type(rankedRhsType.getElementType()))) { + return emitOptionalError( + location, "expects lhs and rhs to have compatible element type. Got: ", + rankedLhsType.getElementType(), " and ", + rankedRhsType.getElementType()); ++ } + + if (failed(verifyConvolutionAttributes( + location, lhsType, rhsType, inputBatchDimension, +@@ -2548,7 +2558,9 @@ + // dynamic_conv_c27 + if (!anyQuantized({rankedLhsType, rankedRhsType}) && + !isCompatibleForHloTypeInference(rankedLhsType.getElementType(), +- rankedRhsType.getElementType())) ++ rankedRhsType.getElementType()) && ++ !(isFp8Type(rankedLhsType.getElementType()) && ++ isFp8Type(rankedRhsType.getElementType()))) + return emitOptionalError( + location, "expects lhs and rhs to have compatible element type. Got: ", + rankedLhsType.getElementType(), " and ", +diff --ruN a/stablehlo/stablehlo/dialect/Version.h b/stablehlo/stablehlo/dialect/Version.h +--- stablehlo/stablehlo/dialect/Version.h ++++ stablehlo/stablehlo/dialect/Version.h +@@ -38,7 +38,7 @@ + static FailureOr fromString(llvm::StringRef versionRef); + + /// Return a Version representing the current VHLO dialect version. +- static Version getCurrentVersion() { return Version(1, 19, 0); } ++ static Version getCurrentVersion() { return Version(1, 20, 0); } + + /// Return a Version representing the minimum supported VHLO dialect version. + static Version getMinimumVersion() { return Version(0, 9, 0); } +diff --ruN a/stablehlo/stablehlo/dialect/VhloDialect.td b/stablehlo/stablehlo/dialect/VhloDialect.td +--- stablehlo/stablehlo/dialect/VhloDialect.td ++++ stablehlo/stablehlo/dialect/VhloDialect.td +@@ -58,6 +58,7 @@ + 1.17.0: Add `future` type support to `custom_call` op. + 1.18.0: Add `result_tilings` attribute to `custom_call` op. + 1.19.0: Add CollectiveReduceOp. ++ 1.20.0: Allow mixed fp8 operands in `convolution` and `dynamic_conv` ops. + }]; + + let useDefaultAttributePrinterParser = 0; +diff --ruN a/stablehlo/stablehlo/dialect/VhloOps.cpp b/stablehlo/stablehlo/dialect/VhloOps.cpp +--- stablehlo/stablehlo/dialect/VhloOps.cpp ++++ stablehlo/stablehlo/dialect/VhloOps.cpp +@@ -359,11 +359,44 @@ + return success(); + } + ++bool isVhloFp8Type(Type type) { ++ return isa(type); ++} ++ ++LogicalResult verifyConstraint_1_20_0(mlir::Operation* op, ++ Version targetVersion) { ++ if (targetVersion < Version(1, 20, 0)) { ++ if (op->getNumOperands() < 2) { ++ return failure(); ++ } ++ Type lhsElementType = getVhloElementType(op->getOperand(0).getType()); ++ Type rhsElementType = getVhloElementType(op->getOperand(1).getType()); ++ if (lhsElementType != rhsElementType && isVhloFp8Type(lhsElementType) && ++ isVhloFp8Type(rhsElementType)) { ++ return failure(); ++ } ++ } ++ return success(); ++} ++ + } // namespace + + LogicalResult AllReduceOpV1::validateConstraint(mlir::Operation* op, + Version targetVersion) { + return verifyConstraint_0_17_0(op, targetVersion); ++} ++ ++LogicalResult ConvolutionOpV1::validateConstraint(mlir::Operation* op, ++ Version targetVersion) { ++ return verifyConstraint_1_20_0(op, targetVersion); ++} ++ ++LogicalResult DynamicConvOpV2::validateConstraint(mlir::Operation* op, ++ Version targetVersion) { ++ return verifyConstraint_1_20_0(op, targetVersion); + } + + LogicalResult ReduceOpV1::validateConstraint(mlir::Operation* op, +diff --ruN a/stablehlo/stablehlo/dialect/VhloOps.td b/stablehlo/stablehlo/dialect/VhloOps.td +--- stablehlo/stablehlo/dialect/VhloOps.td ++++ stablehlo/stablehlo/dialect/VhloOps.td +@@ -390,7 +390,8 @@ + let results = (outs VHLO_AnyType:$result); + } + +-def VHLO_ConvolutionOpV1 : VHLO_Op<"convolution_v1", "0.9.0", "current"> { ++def VHLO_ConvolutionOpV1 : VHLO_Op<"convolution_v1", "0.9.0", "current", ++ [DeclareOpInterfaceMethods]> { + let arguments = (ins + VHLO_AnyType:$lhs, + VHLO_AnyType:$rhs, +@@ -572,7 +573,8 @@ + + // Padding should be specified as an operand only, not an attribute. + // Remove `d_padding` and convert `padding` to an operand. +-def VHLO_DynamicConvOpV2 : VHLO_Op<"dynamic_conv_v2", "0.20.0", "current"> { ++def VHLO_DynamicConvOpV2 : VHLO_Op<"dynamic_conv_v2", "0.20.0", "current", ++ [DeclareOpInterfaceMethods]> { + let arguments = (ins + VHLO_AnyType:$lhs, + VHLO_AnyType:$rhs, diff --ruN a/stablehlo/stablehlo/tests/TestUtils.cpp b/stablehlo/stablehlo/tests/TestUtils.cpp --- stablehlo/stablehlo/tests/TestUtils.cpp +++ stablehlo/stablehlo/tests/TestUtils.cpp @@ -312,6 +488,170 @@ diff --ruN a/stablehlo/stablehlo/tests/ops_broadcasting.mlir b/stablehlo/stableh + return %0 : tensor<3x4x5xf64> +} + +diff --ruN a/stablehlo/stablehlo/tests/verify_convolution.mlir b/stablehlo/stablehlo/tests/verify_convolution.mlir +--- stablehlo/stablehlo/tests/verify_convolution.mlir ++++ stablehlo/stablehlo/tests/verify_convolution.mlir +@@ -78,6 +78,60 @@ + + // ----- + ++// CHECK-LABEL: func @convolution_mixed_fp8 ++func.func @convolution_mixed_fp8(%arg0 : tensor<100x26x26x32xf8E5M2>, ++ %arg1 : tensor<3x3x1x32xf8E4M3FN>) -> tensor<100x28x28x1xf8E5M2> { ++ %result = "stablehlo.convolution"(%arg0, %arg1) { ++ batch_group_count = 1 : i64, ++ dimension_numbers = #stablehlo.conv, ++ feature_group_count = 1 : i64, ++ lhs_dilation = array, ++ padding = dense<2> : tensor<2x2xi64>, ++ rhs_dilation = array, ++ window_strides = array ++ } : (tensor<100x26x26x32xf8E5M2>, tensor<3x3x1x32xf8E4M3FN>) -> ++ tensor<100x28x28x1xf8E5M2> ++ func.return %result : tensor<100x28x28x1xf8E5M2> ++} ++ ++// ----- ++ ++func.func @convolution_mismatched_element_types(%arg0: tensor<100x26x26x32xf32>, %arg1: tensor<3x3x1x32xf16>) -> tensor<100x28x28x1xf32> { ++ // expected-error@+1{{expects lhs and rhs to have compatible element type. Got: 'f32' and 'f16'}} ++ %result = "stablehlo.convolution"(%arg0, %arg1) { ++ batch_group_count = 1 : i64, ++ dimension_numbers = #stablehlo.conv, ++ feature_group_count = 1 : i64, ++ lhs_dilation = array, ++ padding = dense<2> : tensor<2x2xi64>, ++ rhs_dilation = array, ++ window_strides = array ++ } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf16>) -> tensor<100x28x28x1xf32> ++ func.return %result : tensor<100x28x28x1xf32> ++} ++ ++// ----- ++ + func.func @convolution(%arg0: tensor<2x2x3x4xf32>, %arg1: tensor<3x5x5x3xf32>) -> tensor<3x5x5x4xf32> { + // expected-error@+3{{Unexpected keyword stide}} + %0 = stablehlo.convolution(%arg0, %arg1) +diff --ruN a/stablehlo/stablehlo/tests/verify_dynamic_conv.mlir b/stablehlo/stablehlo/tests/verify_dynamic_conv.mlir +--- stablehlo/stablehlo/tests/verify_dynamic_conv.mlir ++++ stablehlo/stablehlo/tests/verify_dynamic_conv.mlir +@@ -48,6 +48,36 @@ + + // ----- + ++// CHECK-LABEL: func @dynamic_conv_mixed_fp8 ++func.func @dynamic_conv_mixed_fp8(%arg0 : tensor<100x26x26x32xf8E5M2>, ++ %arg1 : tensor<3x3x1x32xf8E4M3FN>) -> tensor<100x28x28x1xf8E5M2> { ++ %padding = stablehlo.constant dense<2> : tensor<2x2xi64> ++ %result = "stablehlo.dynamic_conv"(%arg0, %arg1, %padding) { ++ dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>, ++ feature_group_count = 1 : i64, ++ batch_group_count = 1 : i64 ++ } : (tensor<100x26x26x32xf8E5M2>, tensor<3x3x1x32xf8E4M3FN>, tensor<2x2xi64>) -> ++ tensor<100x28x28x1xf8E5M2> ++ func.return %result : tensor<100x28x28x1xf8E5M2> ++} ++ ++// ----- ++ ++func.func @dynamic_conv_mismatched_element_types(%arg0: tensor<100x26x26x32xf32>, ++ %arg1: tensor<3x3x1x32xf16>) -> tensor<100x28x28x1xf32> { ++ // expected-error@+2 {{expects lhs and rhs to have compatible element type. Got: 'f32' and 'f16'}} ++ %padding = stablehlo.constant dense<2> : tensor<2x2xi64> ++ %result = "stablehlo.dynamic_conv"(%arg0, %arg1, %padding) { ++ dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>, ++ feature_group_count = 1 : i64, ++ batch_group_count = 1 : i64 ++ } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf16>, tensor<2x2xi64>) -> ++ tensor<100x28x28x1xf32> ++ func.return %result : tensor<100x28x28x1xf32> ++} ++ ++// ----- ++ + func.func @dynamic_conv_c1(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x207xf32>) -> tensor<1x8x8x16xf32> { + // expected-error@+2 {{expects convolution arguments to have same number of dimensions. Got: 'tensor<1x8x8x207xf32>' and 'tensor<3x3x207xf32>'.}} +diff --ruN a/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_19_0.mlir b/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_19_0.mlir +--- stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_19_0.mlir ++++ stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_19_0.mlir +@@ -0,0 +1,56 @@ ++// RUN: stablehlo-opt --stablehlo-legalize-to-vhlo --vhlo-to-version='target=1.19.0' --verify-diagnostics --split-input-file %s ++ ++// expected-error @+1 {{failed to convert VHLO to v1.19.0}} ++module { ++ func.func @convolution_mixed_fp8(%arg0: tensor<100x26x26x32xf8E5M2>, %arg1: tensor<3x3x1x32xf8E4M3FN>) -> tensor<100x28x28x1xf8E5M2> { ++ // expected-error @+1 {{failed to legalize operation 'vhlo.convolution_v1' that was explicitly marked illegal}} ++ %result = "stablehlo.convolution"(%arg0, %arg1) { ++ batch_group_count = 1 : i64, ++ dimension_numbers = #stablehlo.conv, ++ feature_group_count = 1 : i64, ++ lhs_dilation = array, ++ padding = dense<2> : tensor<2x2xi64>, ++ rhs_dilation = array, ++ window_strides = array ++ } : (tensor<100x26x26x32xf8E5M2>, tensor<3x3x1x32xf8E4M3FN>) -> tensor<100x28x28x1xf8E5M2> ++ func.return %result : tensor<100x28x28x1xf8E5M2> ++ } ++} ++ ++// ----- ++ ++// expected-error @+1 {{failed to convert VHLO to v1.19.0}} ++module { ++ func.func @dynamic_conv_mixed_fp8(%arg0: tensor<100x26x26x32xf8E5M2>, %arg1: tensor<3x3x1x32xf8E4M3FN>, %arg2: tensor<2x2xi64>) -> tensor<100x28x28x1xf8E5M2> { ++ // expected-error @+1 {{failed to legalize operation 'vhlo.dynamic_conv_v2' that was explicitly marked illegal}} ++ %result = "stablehlo.dynamic_conv"(%arg0, %arg1, %arg2) { ++ batch_group_count = 1 : i64, ++ dimension_numbers = #stablehlo.conv, ++ feature_group_count = 1 : i64, ++ lhs_dilation = array, ++ rhs_dilation = array, ++ window_strides = array ++ } : (tensor<100x26x26x32xf8E5M2>, tensor<3x3x1x32xf8E4M3FN>, tensor<2x2xi64>) -> tensor<100x28x28x1xf8E5M2> ++ func.return %result : tensor<100x28x28x1xf8E5M2> ++ } ++} diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp --- stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp +++ stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp diff --git a/third_party/xla/xla/backends/cpu/autotuner/BUILD b/third_party/xla/xla/backends/cpu/autotuner/BUILD index f37dba0ec46e67..460ef37e8b3379 100644 --- a/third_party/xla/xla/backends/cpu/autotuner/BUILD +++ b/third_party/xla/xla/backends/cpu/autotuner/BUILD @@ -142,9 +142,9 @@ xla_cc_test( ) cc_library( - name = "llvm_kernel_autotuner", - srcs = ["llvm_kernel_autotuner.cc"], - hdrs = ["llvm_kernel_autotuner.h"], + name = "cpu_autotuner", + srcs = ["cpu_autotuner.cc"], + hdrs = ["cpu_autotuner.h"], deps = [ ":cpu_codegen_backend", ":cpu_profiler", @@ -171,10 +171,10 @@ cc_library( ) xla_cc_test( - name = "llvm_kernel_autotuner_test", - srcs = ["llvm_kernel_autotuner_test.cc"], + name = "cpu_autotuner_test", + srcs = ["cpu_autotuner_test.cc"], deps = [ - ":llvm_kernel_autotuner", + ":cpu_autotuner", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/service/cpu:cpu_compiler", diff --git a/third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner.cc b/third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner.cc similarity index 92% rename from third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner.cc rename to third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner.cc index 4408a49f4e3ccc..ee7792d167deda 100644 --- a/third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner.cc +++ b/third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/backends/cpu/autotuner/llvm_kernel_autotuner.h" +#include "xla/backends/cpu/autotuner/cpu_autotuner.h" #include #include @@ -42,15 +42,16 @@ limitations under the License. namespace xla::cpu { -absl::StatusOr LlvmKernelAutotuner::RunImpl( +absl::StatusOr CpuAutotuner::RunImpl( HloModule* module, const absl::flat_hash_set& execution_threads) { ABSL_ASSIGN_OR_RETURN(auto compiler, CpuCodegenBackend::CreateBackendCompiler()); - ABSL_ASSIGN_OR_RETURN(auto backend, LlvmKernelBackend::Create(compiler.get())); + ABSL_ASSIGN_OR_RETURN(auto llvm_kernel_backend, + LlvmKernelBackend::Create(compiler.get())); std::unique_ptr profiler = CpuProfiler::Create(ProfileOptions()); std::vector> codegen_backends; - codegen_backends.push_back(std::move(backend)); + codegen_backends.push_back(std::move(llvm_kernel_backend)); ABSL_ASSIGN_OR_RETURN(auto orchestrator, CodegenOrchestrator::Create(std::move(codegen_backends), diff --git a/third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner.h b/third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner.h similarity index 68% rename from third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner.h rename to third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner.h index fe26d971346742..c76882157e527c 100644 --- a/third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner.h +++ b/third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_BACKENDS_CPU_AUTOTUNER_LLVM_KERNEL_AUTOTUNER_H_ -#define XLA_BACKENDS_CPU_AUTOTUNER_LLVM_KERNEL_AUTOTUNER_H_ +#ifndef XLA_BACKENDS_CPU_AUTOTUNER_CPU_AUTOTUNER_H_ +#define XLA_BACKENDS_CPU_AUTOTUNER_CPU_AUTOTUNER_H_ #include "absl/container/flat_hash_set.h" #include "absl/log/check.h" @@ -29,16 +29,14 @@ limitations under the License. namespace xla::cpu { -inline constexpr absl::string_view kLlvmKernelAutotunerName = - "llvm_kernel_autotuner"; +inline constexpr absl::string_view kCpuAutotunerName = "cpu_autotuner"; -// Llvm kernel autotuning pass. It tries to autotune the llvm kernel compilation -// provided by the LlvmKernelBackend. -class LlvmKernelAutotuner : public HloModulePass { +// CPU autotuning pass. +class CpuAutotuner : public HloModulePass { public: - LlvmKernelAutotuner() = default; + CpuAutotuner() = default; - absl::string_view name() const override { return kLlvmKernelAutotunerName; } + absl::string_view name() const override { return kCpuAutotunerName; } protected: absl::StatusOr RunImpl( @@ -48,4 +46,4 @@ class LlvmKernelAutotuner : public HloModulePass { } // namespace xla::cpu -#endif // XLA_BACKENDS_CPU_AUTOTUNER_LLVM_KERNEL_AUTOTUNER_H_ +#endif // XLA_BACKENDS_CPU_AUTOTUNER_CPU_AUTOTUNER_H_ diff --git a/third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner_test.cc b/third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner_test.cc similarity index 86% rename from third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner_test.cc rename to third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner_test.cc index f0b0bee12b7ffe..450bdc89fea5c6 100644 --- a/third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner_test.cc +++ b/third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/backends/cpu/autotuner/llvm_kernel_autotuner.h" +#include "xla/backends/cpu/autotuner/cpu_autotuner.h" #include @@ -36,10 +36,10 @@ constexpr absl::string_view kLlvmKernelConcatenateHlo = R"( } )"; -class LlvmKernelAutotunerTest : public HloHardwareIndependentTestBase {}; +class CpuAutotunerTest : public HloHardwareIndependentTestBase {}; -TEST_F(LlvmKernelAutotunerTest, GetBestConfig) { - LlvmKernelAutotuner autotuner; +TEST_F(CpuAutotunerTest, GetBestConfig) { + CpuAutotuner autotuner; TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr module, ParseAndReturnVerifiedModule(kLlvmKernelConcatenateHlo)); diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD index 76576adcd14f38..6fbd1a47698d7c 100644 --- a/third_party/xla/xla/backends/gpu/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/tests/BUILD @@ -1719,13 +1719,10 @@ xla_cc_test( ], ) -xla_cc_test( +xla_test( name = "command_buffer_profiling_test", srcs = ["command_buffer_profiling_test.cc"], - tags = [ - "cuda-only", - "gpu", - ], + backends = ["gpu"], deps = [ "//xla:xla_proto_cc", "//xla/backends/gpu/runtime:command_buffer_thunk", @@ -1736,11 +1733,11 @@ xla_cc_test( "//xla/service:compiled_module", "//xla/service:compiler", "//xla/service:executable", - "//xla/service:gpu_plugin", "//xla/service:gpu_topology", + "//xla/service:platform_util", "//xla/service/gpu:gpu_executable", - "//xla/stream_executor:device_description_proto_cc", - "//xla/stream_executor/cuda:cuda_platform_id", + "//xla/stream_executor:platform", + "//xla/stream_executor:stream_executor_h", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest_main", ], diff --git a/third_party/xla/xla/backends/gpu/tests/command_buffer_profiling_test.cc b/third_party/xla/xla/backends/gpu/tests/command_buffer_profiling_test.cc index 303d791f1d706f..791999a6c8b0cb 100644 --- a/third_party/xla/xla/backends/gpu/tests/command_buffer_profiling_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/command_buffer_profiling_test.cc @@ -30,8 +30,9 @@ limitations under the License. #include "xla/service/executable.h" #include "xla/service/gpu/gpu_executable.h" #include "xla/service/gpu_topology.h" -#include "xla/stream_executor/cuda/cuda_platform_id.h" -#include "xla/stream_executor/device_description.pb.h" +#include "xla/service/platform_util.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/stream_executor.h" #include "xla/xla.pb.h" namespace xla::gpu { @@ -57,15 +58,19 @@ TEST_P(CommandBufferProfilingTest, ASSERT_OK_AND_ASSIGN(std::unique_ptr hlo_module, ParseAndReturnVerifiedModule(hlo_text)); - ASSERT_OK_AND_ASSIGN( - stream_executor::GpuTargetConfigProto gpu_target_config_proto, - GetGpuTargetConfig(GpuModel::H100_PCIE)); - ASSERT_OK_AND_ASSIGN( - gpu::GpuTargetConfig gpu_target_config, - gpu::GpuTargetConfig::FromProto(gpu_target_config_proto)); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr compiler, - Compiler::GetForPlatform(stream_executor::cuda::kCudaPlatformId)); + ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, + PlatformUtil::GetPlatform("gpu")); + + ASSERT_OK_AND_ASSIGN(std::vector executors, + PlatformUtil::GetStreamExecutors(platform)); + + ASSERT_FALSE(executors.empty()); + + stream_executor::StreamExecutor* executor = executors.front(); + GpuTargetConfig gpu_target_config(executor); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, + Compiler::GetForPlatform(platform->id())); AotCompilationOptions aot_options(compiler->PlatformId()); aot_options.set_gpu_topology( diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index 0782615d4ddab4..3d26b3168889e9 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -1734,10 +1734,7 @@ cc_library( xla_cc_test( name = "dynamic_slice_fusion_rewriter_v2_test", srcs = ["dynamic_slice_fusion_rewriter_v2_test.cc"], - tags = [ - "cuda-only", - "gpu", - ], + tags = ["gpu"], deps = [ ":dynamic_slice_annotator", ":dynamic_slice_fusion", @@ -3835,7 +3832,6 @@ cc_library( xla_test( name = "triton_fusion_numerics_verifier_test", - timeout = "short", srcs = ["triton_fusion_numerics_verifier_test.cc"], backends = [ "a100", diff --git a/third_party/xla/xla/backends/profiler/gpu/rocm_collector.h b/third_party/xla/xla/backends/profiler/gpu/rocm_collector.h index cafdc0bde36591..602e25f8302eb0 100644 --- a/third_party/xla/xla/backends/profiler/gpu/rocm_collector.h +++ b/third_party/xla/xla/backends/profiler/gpu/rocm_collector.h @@ -50,8 +50,16 @@ inline std::string ToXStat(const KernelDetails& kernel_info, grid_z = kernel_info.workgroup_z != 0 ? kernel_info.grid_z / kernel_info.workgroup_z : 0; - - return absl::StrCat(" grid:", grid_x, ",", grid_y, ",", grid_z, + const uint32_t dynamic_group_segment_size = + kernel_info.group_segment_size > kernel_info.static_group_segment_size + ? kernel_info.group_segment_size - + kernel_info.static_group_segment_size + : 0; + + return absl::StrCat("regs:", kernel_info.registers_per_work_item, + " static_shared:", kernel_info.static_group_segment_size, + " dynamic_shared:", dynamic_group_segment_size, + " grid:", grid_x, ",", grid_y, ",", grid_z, " block:", kernel_info.workgroup_x, ",", kernel_info.workgroup_y, ",", kernel_info.workgroup_z, " private_mem:", kernel_info.private_segment_size, diff --git a/third_party/xla/xla/backends/profiler/gpu/rocm_collector_test.cc b/third_party/xla/xla/backends/profiler/gpu/rocm_collector_test.cc index 51cba12620f24c..65cd33905a1298 100644 --- a/third_party/xla/xla/backends/profiler/gpu/rocm_collector_test.cc +++ b/third_party/xla/xla/backends/profiler/gpu/rocm_collector_test.cc @@ -36,6 +36,21 @@ namespace test { using tsl::profiler::FindOrAddMutablePlaneWithName; using tsl::profiler::XSpace; +TEST(RocmCollectorTest, ToXStatDecomposesDispatchGroupMemory) { + KernelDetails kernel_info{}; + kernel_info.group_segment_size = 1536; + kernel_info.static_group_segment_size = 1024; + + EXPECT_NE(ToXStat(kernel_info, /*occupancy_pct=*/0) + .find("static_shared:1024 dynamic_shared:512"), + std::string::npos); + + kernel_info.group_segment_size = 512; + EXPECT_NE(ToXStat(kernel_info, /*occupancy_pct=*/0) + .find("static_shared:1024 dynamic_shared:0"), + std::string::npos); +} + TEST(RocmCollectorTest, TestAddKernelEventAndExport) { RocmTraceCollectorOptions options; options.max_callback_api_events = 100; diff --git a/third_party/xla/xla/backends/profiler/gpu/rocm_tracer.cc b/third_party/xla/xla/backends/profiler/gpu/rocm_tracer.cc index f9c37cfac6cd1d..5a47fb709fb7e7 100644 --- a/third_party/xla/xla/backends/profiler/gpu/rocm_tracer.cc +++ b/third_party/xla/xla/backends/profiler/gpu/rocm_tracer.cc @@ -361,7 +361,13 @@ void RocmTracer::KernelEvent(const rocprofiler_record_header_t* hdr, }; auto it = kernel_info_.find(kinfo.kernel_id); - if (it != kernel_info_.end()) trace_event->name = it->second.name; + if (it != kernel_info_.end()) { + trace_event->name = it->second.name; + const auto& sym = it->second.data; + trace_event->kernel_info.registers_per_work_item = + sym.arch_vgpr_count + sym.accum_vgpr_count; + trace_event->kernel_info.static_group_segment_size = sym.group_segment_size; + } } void RocmTracer::EmitMarkerEvent(std::string label, uint64_t start_ns, diff --git a/third_party/xla/xla/backends/profiler/gpu/rocm_tracer_utils.h b/third_party/xla/xla/backends/profiler/gpu/rocm_tracer_utils.h index 7c93538f9ce640..0de4064154dfe0 100644 --- a/third_party/xla/xla/backends/profiler/gpu/rocm_tracer_utils.h +++ b/third_party/xla/xla/backends/profiler/gpu/rocm_tracer_utils.h @@ -59,11 +59,15 @@ struct MemsetDetails { }; struct KernelDetails { - // The amount of private memory used by kernel, - // number of register per thread (register spillage if > 0) + // Total dispatch-time private-segment (scratch) bytes per work-item. uint32_t private_segment_size; - // The amount of shared memory (SMEM) + // Total dispatch-time group-segment (LDS) bytes per workgroup. Includes + // static and dynamic LDS allocation. uint32_t group_segment_size; + // Architecture and accumulator VGPRs allocated per work-item. + uint32_t registers_per_work_item; + // Static group-segment (LDS) bytes per workgroup from the kernel symbol. + uint32_t static_group_segment_size; // X-dimension of a workgroup (grid.x*block.x) uint32_t workgroup_x; // Y-dimension of a workgroup (grid.x*block.x) diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc index e148e25c1065e4..cb64f9eb0defb3 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc @@ -2143,9 +2143,12 @@ absl::StatusOr HloFunctionImporter::ImportInstructionImpl( auto lhs_element_type = instruction->operand(0)->shape().element_type(); auto rhs_element_type = instruction->operand(1)->shape().element_type(); if (lhs_element_type != rhs_element_type) { - // Cast LHS or RHS to the common element type. - if (primitive_util::CastPreservesValues(lhs_element_type, - rhs_element_type)) { + if (primitive_util::IsF8Type(lhs_element_type) && + primitive_util::IsF8Type(rhs_element_type)) { + // Allow mixed FP8 without conversion. + } else if (primitive_util::CastPreservesValues(lhs_element_type, + rhs_element_type)) { + // Cast LHS to the common element type. auto convert_op_return_type = mlir::cast(lhs.getType()) .clone(mlir::getElementTypeOrSelf(rhs)); @@ -2153,6 +2156,7 @@ absl::StatusOr HloFunctionImporter::ImportInstructionImpl( convert_op_return_type, lhs); } else if (primitive_util::CastPreservesValues(rhs_element_type, lhs_element_type)) { + // Cast RHS to the common element type. auto convert_op_return_type = mlir::cast(rhs.getType()) .clone(mlir::getElementTypeOrSelf(lhs)); diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu.td b/third_party/xla/xla/mosaic/dialect/tpu/tpu.td index 0231285f239bc2..d4e241c0f16f3e 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu.td +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu.td @@ -25,12 +25,14 @@ def TPU_Dialect : Dialect { let useDefaultAttributePrinterParser = 1; let useDefaultTypePrinterParser = 1; let extraClassDeclaration = [{ - static StringRef GetCoreTypeKey() { return "tpu.core_type"; } + static constexpr StringRef GetMainKey() { return "tpu.main"; } + static constexpr StringRef GetCoreTypeKey() { return "tpu.core_type"; } static std::optional GetCoreTypeAttr(Operation *op); }]; let hasConstantMaterializer = 1; let hasCanonicalizer = 1; + let hasOperationAttrVerify = 1; } class TPU_Attr traits = []> diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.cc index 0a9e036b37033e..3064f345c4710b 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.cc @@ -43,6 +43,7 @@ limitations under the License. #include "mlir/Dialect/Utils/StaticValueUtils.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineMap.h" +#include "mlir/IR/Attributes.h" #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/BuiltinTypeInterfaces.h" @@ -115,6 +116,19 @@ Operation* TPUDialect::materializeConstant(OpBuilder& builder, Attribute value, return mlir::cast(attr).getValue(); } +LogicalResult TPUDialect::verifyOperationAttribute(Operation* const op, + NamedAttribute attr) { + CHECK(op != nullptr); + if (attr.getName() == GetMainKey()) { + if (!isa(op) || !isa(attr.getValue())) { + return op->emitOpError() << GetMainKey() + << " attribute is expected to be a unit " + "attribute on a func.func operation"; + } + } + return success(); +} + // Rewrites // // memref.dim(tpu.memref_slice(..., dynamic_sizes), i) @@ -216,9 +230,27 @@ CoreType GetCoreTypeOfParentOp(Operation& op) { return parent ? *TPUDialect::GetCoreTypeAttr(parent) : CoreType::kTc; } -absl::StatusOr GetFuncWithCoreType(ModuleOp module, - CoreType core_type) { +absl::StatusOr GetMainFunc(ModuleOp module, CoreType core_type) { func::FuncOp result = nullptr; + for (auto func_op : module.getOps()) { + if (!func_op->hasAttr(TPUDialect::GetMainKey())) { + continue; + } + if (TPUDialect::GetCoreTypeAttr(func_op) != core_type) { + continue; + } + if (result != nullptr) { + return absl::InvalidArgumentError(absl::StrFormat( + "Multiple functions with %v attribute and tpu.core_type = %v found", + TPUDialect::GetMainKey(), core_type)); + } + result = func_op; + } + if (result != nullptr) { + return result; + } + // If there is no function marked with tpu.main, fall back to finding a + // (unique) function with the requested core type. for (func::FuncOp func_op : module.getOps()) { if (TPUDialect::GetCoreTypeAttr(func_op) != core_type) { continue; diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h b/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h index a488ee1d7c98b5..3f886e4a1a2347 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h @@ -84,9 +84,8 @@ Operation* GetParentOpWithCoreType(Operation& op); // annotation is found, returns kTc. CoreType GetCoreTypeOfParentOp(Operation& op); -// Returns the function in the module with the given core type. -absl::StatusOr GetFuncWithCoreType(ModuleOp module, - CoreType core_type); +// Returns the main function in the module for the given core type. +absl::StatusOr GetMainFunc(ModuleOp module, CoreType core_type); // Changes the memory space of the value and propagates it through the program. LogicalResult specializeMemorySpace(TypedValue value, diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc index 9a1f98f70f82ef..387b22917dd284 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc @@ -337,9 +337,8 @@ StreamExecutorGpuTopologyDescription::GetDefaultDeviceAssignment( stream_executor::PlatformId se_platform_id, StreamExecutorPlatformIdMapping::Global().GetStreamExecutorPlatformId( platform_id())); - ABSL_ASSIGN_OR_RETURN(auto* placer, - ComputationPlacer::GetForPlatform(se_platform_id)); - return placer->AssignDevices(num_replicas, num_partitions); + return ComputationPlacer::GetForPlatform(se_platform_id) + ->AssignDevices(num_replicas, num_partitions); } absl::StatusOr diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index bab0c3b7b4031a..ec981df50f3178 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -3676,19 +3676,13 @@ cc_library( hdrs = ["computation_placer.h"], deps = [ ":device_assignment", - "//xla:util", "//xla/stream_executor:platform_id", - "//xla/stream_executor/cuda:cuda_platform_id", - "//xla/stream_executor/host:host_platform_id", - "//xla/stream_executor/rocm:rocm_platform_id", - "//xla/stream_executor/sycl:sycl_platform_id", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/synchronization", ], - alwayslink = True, # Contains per-platform computation placer registration ) xla_cc_test( @@ -3697,6 +3691,8 @@ xla_cc_test( deps = [ ":computation_placer", ":device_assignment", + "//xla/stream_executor:platform_id", + "@com_google_absl//absl/status:statusor", "@com_google_googletest//:gtest_main", ], ) diff --git a/third_party/xla/xla/service/backend.cc b/third_party/xla/xla/service/backend.cc index 8f3e10fdb268fb..016e295a2316f1 100644 --- a/third_party/xla/xla/service/backend.cc +++ b/third_party/xla/xla/service/backend.cc @@ -190,8 +190,8 @@ CreateGpuAllocators(const se::Platform* platform, PlatformUtil::GetStreamExecutors(platform, options.allowed_devices())); ABSL_ASSIGN_OR_RETURN(auto transfer_manager, TransferManager::GetForPlatform(platform)); - ABSL_ASSIGN_OR_RETURN(auto computation_placer, - ComputationPlacer::GetForPlatform(platform->id())); + ComputationPlacer* computation_placer = + ComputationPlacer::GetForPlatform(platform->id()); std::unique_ptr backend(new Backend( platform, std::move(compiler), stream_executors, transfer_manager, computation_placer, options.intra_op_parallelism_threads())); diff --git a/third_party/xla/xla/service/computation_placer.cc b/third_party/xla/xla/service/computation_placer.cc index db8bce900e5290..82730dfca92136 100644 --- a/third_party/xla/xla/service/computation_placer.cc +++ b/third_party/xla/xla/service/computation_placer.cc @@ -24,12 +24,7 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/synchronization/mutex.h" #include "xla/service/device_assignment.h" -#include "xla/stream_executor/cuda/cuda_platform_id.h" -#include "xla/stream_executor/host/host_platform_id.h" #include "xla/stream_executor/platform_id.h" -#include "xla/stream_executor/rocm/rocm_platform_id.h" -#include "xla/stream_executor/sycl/sycl_platform_id.h" -#include "xla/util.h" namespace xla { @@ -62,6 +57,11 @@ PlacerFactoryMap& GetPlatformComputationPlacers() { static PlacerFactoryMap* const r = new PlacerFactoryMap; return *r; } + +ComputationPlacer* GetDefaultComputationPlacer() { + static auto* const default_placer = new ComputationPlacer; + return default_placer; +} } // namespace /* static */ @@ -77,16 +77,14 @@ void ComputationPlacer::RegisterComputationPlacer( } /* static */ -absl::StatusOr ComputationPlacer::GetForPlatform( +ComputationPlacer* ComputationPlacer::GetForPlatform( se::PlatformId platform_id) { absl::MutexLock lock(placer_mutex); PlacerFactoryMap& placers = GetPlatformComputationPlacers(); auto it = placers.find(platform_id); if (it == placers.end()) { - return NotFound( - "Could not find registered computation placer for platform %s", - platform_id->ToName()); + return GetDefaultComputationPlacer(); } PlacerState& state = it->second; @@ -94,28 +92,7 @@ absl::StatusOr ComputationPlacer::GetForPlatform( // Lazily create the computation placer the first time it is needed. state.placer = state.creation_function(); } - return state.placer.get(); + return state.placer ? state.placer.get() : GetDefaultComputationPlacer(); } } // namespace xla - -namespace { -// registering default computation placer factory for common platforms. -std::unique_ptr DefaultComputationPlacer() { - return std::make_unique(); -} - -bool InitModule() { - xla::ComputationPlacer::RegisterComputationPlacer( - stream_executor::host::kHostPlatformId, DefaultComputationPlacer); - xla::ComputationPlacer::RegisterComputationPlacer( - stream_executor::cuda::kCudaPlatformId, DefaultComputationPlacer); - xla::ComputationPlacer::RegisterComputationPlacer( - stream_executor::rocm::kROCmPlatformId, DefaultComputationPlacer); - xla::ComputationPlacer::RegisterComputationPlacer( - stream_executor::sycl::kSyclPlatformId, DefaultComputationPlacer); - return true; -} - -bool module_initialized = InitModule(); -} // namespace diff --git a/third_party/xla/xla/service/computation_placer.h b/third_party/xla/xla/service/computation_placer.h index 84ed30aaf49289..2f0bbe33d55679 100644 --- a/third_party/xla/xla/service/computation_placer.h +++ b/third_party/xla/xla/service/computation_placer.h @@ -43,10 +43,9 @@ class ComputationPlacer { static void RegisterComputationPlacer(se::PlatformId platform_id, CreationFunction creation_function); - // Returns the computation placer singleton pointer if it is available for the - // given platform, or an error status if it is not. - static absl::StatusOr GetForPlatform( - se::PlatformId platform_id); + // Returns the computation placer singleton pointer registered for the given + // platform, or the default computation placer if none is registered. + static ComputationPlacer* GetForPlatform(se::PlatformId platform_id); private: ComputationPlacer(const ComputationPlacer&) = delete; diff --git a/third_party/xla/xla/service/computation_placer_test.cc b/third_party/xla/xla/service/computation_placer_test.cc index b30b6b6e39aa85..acce0bc0a02bd5 100644 --- a/third_party/xla/xla/service/computation_placer_test.cc +++ b/third_party/xla/xla/service/computation_placer_test.cc @@ -15,9 +15,16 @@ limitations under the License. #include "xla/service/computation_placer.h" +#include + #include #include +#include "absl/status/statusor.h" #include "xla/service/device_assignment.h" +#include "xla/stream_executor/platform_id.h" + +PLATFORM_DEFINE_ID(kUnregisteredPlatformId, UnregisteredPlatform); +PLATFORM_DEFINE_ID(kCustomPlatformId, CustomPlatform); namespace xla { namespace { @@ -33,5 +40,37 @@ TEST(ComputationPlacerTest, Basic) { EXPECT_EQ(da(0, 1), 4); } +TEST(ComputationPlacerTest, GetForPlatformUnregisteredReturnsDefaultPlacer) { + ComputationPlacer* placer = + ComputationPlacer::GetForPlatform(kUnregisteredPlatformId); + ASSERT_NE(placer, nullptr); + ASSERT_OK_AND_ASSIGN(DeviceAssignment da, placer->AssignDevices(2, 2)); + EXPECT_EQ(da(0, 0), 0); + EXPECT_EQ(da(1, 0), 1); + EXPECT_EQ(da(0, 1), 2); + EXPECT_EQ(da(1, 1), 3); +} + +TEST(ComputationPlacerTest, GetForPlatformRegisteredReturnsCustomPlacer) { + class CustomPlacer : public ComputationPlacer { + public: + absl::StatusOr AssignDevices( + int replica_count, int computation_count) override { + DeviceAssignment assignment(replica_count, computation_count); + assignment.Fill(42); + return assignment; + } + }; + + ComputationPlacer::RegisterComputationPlacer( + kCustomPlatformId, []() { return std::make_unique(); }); + + ComputationPlacer* placer = + ComputationPlacer::GetForPlatform(kCustomPlatformId); + ASSERT_NE(placer, nullptr); + ASSERT_OK_AND_ASSIGN(DeviceAssignment da, placer->AssignDevices(1, 1)); + EXPECT_EQ(da(0, 0), 42); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD index 59b301065844fb..99dc48e2494892 100644 --- a/third_party/xla/xla/service/cpu/BUILD +++ b/third_party/xla/xla/service/cpu/BUILD @@ -1636,6 +1636,7 @@ cc_library( hdrs = ["cpu_float_support.h"], copts = tsl_copts(), deps = [ + "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/service:float_support", ], @@ -1648,14 +1649,11 @@ xla_cc_test( ":cpu_float_support", "//xla:shape_util", "//xla:xla_data_proto_cc", - "//xla/backends/cpu/codegen:target_machine_features", "//xla/backends/cpu/codegen:target_machine_test_base", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:verified_hlo_module", "//xla/hlo/transforms/simplifiers:float_normalization", "//xla/service:hlo_module_config", - "//xla/tsl/platform:statusor", - "//xla/tsl/platform:test", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest_main", diff --git a/third_party/xla/xla/service/cpu/cpu_float_support.h b/third_party/xla/xla/service/cpu/cpu_float_support.h index 2a5ea4670bc31f..e9502839bb008a 100644 --- a/third_party/xla/xla/service/cpu/cpu_float_support.h +++ b/third_party/xla/xla/service/cpu/cpu_float_support.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef XLA_SERVICE_CPU_CPU_FLOAT_SUPPORT_H_ #define XLA_SERVICE_CPU_CPU_FLOAT_SUPPORT_H_ +#include #include #include "xla/hlo/ir/hlo_casting_utils.h" @@ -23,6 +24,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/service/float_support.h" +#include "xla/xla_data.pb.h" namespace xla { namespace cpu { @@ -36,6 +38,41 @@ class CpuFloatSupport : public FloatSupport { : FloatSupport(low_precision_type), call_library_for_instruction_(call_library_for_instruction) {} + bool SupportsLowPrecisionOperand(const HloInstruction& hlo, + int64_t operand_index) const override { + if (LowPrecisionType() == BF16 || LowPrecisionType() == F16) { + switch (hlo.opcode()) { + case HloOpcode::kSort: + case HloOpcode::kCompare: + case HloOpcode::kSelect: + return true; + default: + break; + } + } + return FloatSupport::SupportsLowPrecisionOperand(hlo, operand_index); + } + + bool SupportsLowPrecisionOutput(const HloInstruction& hlo) const override { + if (LowPrecisionType() == BF16 || LowPrecisionType() == F16) { + switch (hlo.opcode()) { + case HloOpcode::kSort: + case HloOpcode::kSelect: + return true; + default: + break; + } + } + return FloatSupport::SupportsLowPrecisionOutput(hlo); + } + + bool SupportsMixedPrecisions(const HloInstruction& hlo) const override { + if (hlo.opcode() == HloOpcode::kSort) { + return true; + } + return FloatSupport::SupportsMixedPrecisions(hlo); + } + // Skip trying to upcast the dot if the dot is supported by a library. bool ShouldSkipInstruction(const HloInstruction& hlo) const override { return (hlo.opcode() == HloOpcode::kDot || diff --git a/third_party/xla/xla/service/cpu/cpu_float_support_test.cc b/third_party/xla/xla/service/cpu/cpu_float_support_test.cc index 9008d97488be48..1891860619ff0f 100644 --- a/third_party/xla/xla/service/cpu/cpu_float_support_test.cc +++ b/third_party/xla/xla/service/cpu/cpu_float_support_test.cc @@ -17,14 +17,13 @@ limitations under the License. #include #include -#include #include -#include "absl/strings/match.h" +#include +#include #include "absl/strings/str_cat.h" #include "absl/strings/str_replace.h" #include "absl/strings/string_view.h" -#include "xla/backends/cpu/codegen/target_machine_features.h" #include "xla/backends/cpu/codegen/target_machine_test_base.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" @@ -33,8 +32,6 @@ limitations under the License. #include "xla/service/hlo_module_config.h" #include "xla/shape.h" #include "xla/shape_util.h" -#include "xla/tsl/platform/statusor.h" -#include "xla/tsl/platform/test.h" #include "xla/xla_data.pb.h" namespace xla::cpu { @@ -153,7 +150,7 @@ TEST_P(SkipInstructionTest, Bf16InF32Out) { // Run FloatNormalization and check the results. FloatNormalization float_normalization(&cpu_float_support); - TF_ASSERT_OK_AND_ASSIGN(bool upcast, float_normalization.Run(module.get())); + ASSERT_OK_AND_ASSIGN(bool upcast, float_normalization.Run(module.get())); EXPECT_EQ(upcast, spec.upcast); PrimitiveType expected_input_dtype = spec.upcast ? F32 : BF16; CheckDtype(module.get(), expected_input_dtype, expected_input_dtype, F32); @@ -202,5 +199,123 @@ INSTANTIATE_TEST_SUITE_P(SkipInstructionTestSuite, SkipInstructionTest, ::testing::ValuesIn(GetSkipInstructionTestSpecs()), SkipInstructionTest::Name); +TEST_F(TargetMachineTestBase, SortNotUpcast) { + for (absl::string_view type_str : {"bf16", "f16"}) { + std::string hlo_text = absl::StrReplaceAll(R"( +HloModule test_module + +compare { + p0 = $type$[] parameter(0) + p1 = $type$[] parameter(1) + p2 = s32[] parameter(2) + p3 = s32[] parameter(3) + ROOT cmp = pred[] compare(p0, p1), direction=LT +} + +ENTRY main { + k = $type$[100] parameter(0) + v = s32[100] parameter(1) + ROOT sort = ($type$[100], s32[100]) sort(k, v), dimensions={0}, to_apply=compare, is_stable=true +} +)", + {{"$type$", type_str}}); + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); + + PrimitiveType low_precision_type = (type_str == "bf16") ? BF16 : F16; + + CpuFloatSupport cpu_float_support( + low_precision_type, [](const HloInstruction&) { return false; }); + + FloatNormalization float_normalization(&cpu_float_support); + ASSERT_OK_AND_ASSIGN(bool upcast, float_normalization.Run(module.get())); + EXPECT_FALSE(upcast); + + HloInstruction* root = module->entry_computation()->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kSort); + EXPECT_EQ(root->operand(0)->shape().element_type(), low_precision_type); + EXPECT_EQ(root->operand(1)->shape().element_type(), S32); + + HloComputation* compare_comp = root->to_apply(); + EXPECT_EQ(compare_comp->parameter_instruction(0)->shape().element_type(), + low_precision_type); + } +} + +TEST_F(TargetMachineTestBase, SortUpcastForUnsupportedType) { + std::string hlo_text = R"( +HloModule test_module + +compare { + p0 = f8e5m2[] parameter(0) + p1 = f8e5m2[] parameter(1) + p2 = s32[] parameter(2) + p3 = s32[] parameter(3) + ROOT cmp = pred[] compare(p0, p1), direction=LT +} + +ENTRY main { + k = f8e5m2[100] parameter(0) + v = s32[100] parameter(1) + ROOT sort = (f8e5m2[100], s32[100]) sort(k, v), dimensions={0}, to_apply=compare, is_stable=true +} +)"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); + + CpuFloatSupport cpu_float_support( + F8E5M2, [](const HloInstruction&) { return false; }); + + FloatNormalization float_normalization(&cpu_float_support); + ASSERT_OK_AND_ASSIGN(bool upcast, float_normalization.Run(module.get())); + EXPECT_TRUE(upcast); + + HloInstruction* sort_instr = FindInstruction(module.get(), HloOpcode::kSort); + ASSERT_NE(sort_instr, nullptr); + EXPECT_EQ(sort_instr->operand(0)->shape().element_type(), F32); + EXPECT_EQ(sort_instr->operand(1)->shape().element_type(), S32); + HloComputation* updated_cmp = sort_instr->to_apply(); + EXPECT_EQ(updated_cmp->parameter_instruction(0)->shape().element_type(), F32); + EXPECT_EQ(updated_cmp->parameter_instruction(1)->shape().element_type(), F32); +} + +TEST_F(TargetMachineTestBase, CompareAndSelectNotUpcast) { + for (absl::string_view type_str : {"bf16", "f16"}) { + std::string hlo_text = absl::StrReplaceAll(R"( +HloModule test_module + +ENTRY main { + p0 = $type$[100] parameter(0) + p1 = $type$[100] parameter(1) + cmp = pred[100] compare(p0, p1), direction=LT + ROOT select = $type$[100] select(cmp, p0, p1) +} +)", + {{"$type$", type_str}}); + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); + + PrimitiveType low_precision_type = (type_str == "bf16") ? BF16 : F16; + + CpuFloatSupport cpu_float_support( + low_precision_type, [](const HloInstruction&) { return false; }); + + FloatNormalization float_normalization(&cpu_float_support); + ASSERT_OK_AND_ASSIGN(bool upcast, float_normalization.Run(module.get())); + EXPECT_FALSE(upcast); + + HloInstruction* root = module->entry_computation()->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kSelect); + EXPECT_EQ(root->shape().element_type(), low_precision_type); + EXPECT_EQ(root->operand(1)->shape().element_type(), low_precision_type); + EXPECT_EQ(root->operand(2)->shape().element_type(), low_precision_type); + + const HloInstruction* cmp = root->operand(0); + EXPECT_EQ(cmp->opcode(), HloOpcode::kCompare); + EXPECT_EQ(cmp->operand(0)->shape().element_type(), low_precision_type); + EXPECT_EQ(cmp->operand(1)->shape().element_type(), low_precision_type); + } +} + } // namespace } // namespace xla::cpu diff --git a/third_party/xla/xla/service/cpu/onednn_float_support.cc b/third_party/xla/xla/service/cpu/onednn_float_support.cc index 84495f0def5b02..b749cccbdd49f6 100644 --- a/third_party/xla/xla/service/cpu/onednn_float_support.cc +++ b/third_party/xla/xla/service/cpu/onednn_float_support.cc @@ -50,6 +50,7 @@ bool OneDnnFloatSupport::IsSupported(const HloInstruction& hlo) const { case HloOpcode::kReshape: case HloOpcode::kReverse: case HloOpcode::kScatter: + case HloOpcode::kCompare: case HloOpcode::kSelect: case HloOpcode::kSelectAndScatter: case HloOpcode::kSlice: @@ -57,6 +58,8 @@ bool OneDnnFloatSupport::IsSupported(const HloInstruction& hlo) const { // Other special ops. case HloOpcode::kBitcast: return true; + case HloOpcode::kSort: + return LowPrecisionType() == BF16 || LowPrecisionType() == F16; default: return false; } diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index eb3e044fe4fc25..bbe83ba91d3b30 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -2655,16 +2655,14 @@ xla_test( ], ) -xla_cc_test( +xla_test( name = "gpu_aot_compilation_test", srcs = [ "gpu_aot_compilation_test.cc", ], + backends = ["gpu"], tags = [ - "gpu", - "no_oss", "nomsan", # Pulls in precompiled NVIDIA libraries which cause false positives in msan. - "requires-gpu-nvidia", ], deps = if_cuda_is_configured([ ":nvptx_compiler_impl", @@ -2678,26 +2676,24 @@ xla_cc_test( "//xla:literal_util", "//xla:xla_proto_cc", "//xla/backends/gpu/codegen/triton:support", + "//xla/backends/gpu/tests:hlo_pjrt_gpu_test_base", "//xla/hlo/ir:hlo", + "//xla/pjrt/proto:compile_options_proto_cc", "//xla/service:compiled_module", "//xla/service:compiler", "//xla/service:executable", - "//xla/service:gpu_plugin", "//xla/service:gpu_topology", "//xla/service:hlo_runner_interface", - "//xla/service:platform_util", - "//xla/stream_executor:platform", - "//xla/stream_executor:platform_manager", - "//xla/stream_executor:stream_executor_h", "//xla/tests:literal_test_util", - "//xla/tests:xla_internal_test_main", # build_cleaner: keep - "//xla/tests/restricted:hlo_test_base_legacy", "//xla/tsl/platform:statusor", + "//xla/util/split_proto:split_executable_and_options_writer", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", "@com_google_googletest//:gtest", + "@com_google_googletest//:gtest_main", "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", + "@riegeli//riegeli/bytes:string_writer", "@tsl//tsl/platform:statusor", ], ) diff --git a/third_party/xla/xla/service/gpu/gpu_aot_compilation_test.cc b/third_party/xla/xla/service/gpu/gpu_aot_compilation_test.cc index efff25b0ff66d1..3c6f46c29f779f 100644 --- a/third_party/xla/xla/service/gpu/gpu_aot_compilation_test.cc +++ b/third_party/xla/xla/service/gpu/gpu_aot_compilation_test.cc @@ -20,7 +20,6 @@ limitations under the License. #include #include -#include "absl/strings/ascii.h" #include "absl/strings/escaping.h" #include "absl/strings/str_format.h" #include "absl/strings/string_view.h" @@ -30,22 +29,21 @@ limitations under the License. #include "mlir/IR/Builders.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/MLIRContext.h" +#include "riegeli/bytes/string_writer.h" #include "xla/backends/gpu/codegen/triton/support.h" +#include "xla/backends/gpu/tests/hlo_pjrt_gpu_test_base.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/literal.h" #include "xla/literal_util.h" +#include "xla/pjrt/proto/compile_options.pb.h" #include "xla/service/compiled_module.h" #include "xla/service/compiler.h" #include "xla/service/executable.h" #include "xla/service/gpu/gpu_executable.h" #include "xla/service/gpu_topology.h" #include "xla/service/hlo_runner_interface.h" -#include "xla/service/platform_util.h" -#include "xla/stream_executor/platform.h" -#include "xla/stream_executor/platform_manager.h" -#include "xla/stream_executor/stream_executor.h" #include "xla/tests/literal_test_util.h" -#include "xla/tests/restricted/hlo_test_base_legacy.h" +#include "xla/util/split_proto/split_executable_and_options_writer.h" #include "xla/xla.pb.h" namespace xla { @@ -53,9 +51,12 @@ namespace gpu { using ::testing::IsEmpty; using ::testing::Not; -class GpuAotCompilationTest : public HloTestBaseLegacy { +class GpuAotCompilationTest : public HloPjRtGpuTestBase { protected: - void SetUp() override { debug_options_ = GetDebugOptionsForTest(); } + void SetUp() override { + HloPjRtGpuTestBase::SetUp(); + debug_options_ = GetDebugOptionsForTest(); + } DebugOptions debug_options_; }; @@ -72,35 +73,28 @@ TEST_F(GpuAotCompilationTest, ExportAndLoadExecutable) { ASSERT_OK_AND_ASSIGN(std::unique_ptr module, ParseAndReturnVerifiedModule(hlo_string)); - auto compiler = backend().compiler(); - auto name = - absl::AsciiStrToUpper(PlatformUtil::CanonicalPlatformName("gpu").value()); - ASSERT_OK_AND_ASSIGN(se::Platform * platform, - se::PlatformManager::PlatformWithName(name)); - ASSERT_OK_AND_ASSIGN(se::StreamExecutor * stream_exec, - platform->ExecutorForDevice(0)); - // Compile AOT. - AotCompilationOptions aot_options(compiler->PlatformId()); - aot_options.set_executor(stream_exec); + AotCompilationOptions aot_options(stream_executor_platform_id()); + aot_options.set_gpu_topology( + GetSingleDeviceGpuTopology("", gpu_target_config())); ASSERT_OK_AND_ASSIGN( std::vector> aot_results, - compiler->CompileAheadOfTime(std::move(module), aot_options)); + compiler()->CompileAheadOfTime(std::move(module), aot_options)); // Serialize-deserialize AOT compilation result. ASSERT_OK_AND_ASSIGN(std::string serialized_aot_result, aot_results[0]->SerializeAsString()); ASSERT_OK_AND_ASSIGN( std::unique_ptr aot_result, - compiler->LoadAotCompilationResult(serialized_aot_result)); + compiler()->LoadAotCompilationResult(serialized_aot_result)); // Load Executable from AOT compilation result. ASSERT_OK_AND_ASSIGN( std::unique_ptr executable, std::move(*aot_result) - .LoadExecutable(compiler->PlatformId(), - stream_exec->GetDeviceDescription(), debug_options_)); + .LoadExecutable(stream_executor_platform_id(), device_description(), + debug_options_)); auto* gpu_executable = dynamic_cast(executable.get()); ASSERT_NE(gpu_executable, nullptr); @@ -120,37 +114,28 @@ TEST_F(GpuAotCompilationTest, AotCompilationWithoutGpuDevice) { ASSERT_OK_AND_ASSIGN(std::unique_ptr module, ParseAndReturnVerifiedModule(hlo_string)); - auto compiler = backend().compiler(); - auto name = - absl::AsciiStrToUpper(PlatformUtil::CanonicalPlatformName("gpu").value()); - ASSERT_OK_AND_ASSIGN(se::Platform * platform, - se::PlatformManager::PlatformWithName(name)); - ASSERT_OK_AND_ASSIGN(se::StreamExecutor * stream_exec, - platform->ExecutorForDevice(0)); - // Stream executor is not passed as an option. - Compiler::GpuTargetConfig gpu_target_config(stream_exec); - AotCompilationOptions aot_options(compiler->PlatformId()); + AotCompilationOptions aot_options(stream_executor_platform_id()); aot_options.set_gpu_topology( - GetSingleDeviceGpuTopology("", gpu_target_config)); + GetSingleDeviceGpuTopology("", gpu_target_config())); ASSERT_OK_AND_ASSIGN( std::vector> aot_results, - compiler->CompileAheadOfTime(std::move(module), aot_options)); + compiler()->CompileAheadOfTime(std::move(module), aot_options)); // Serialize-deserialize AOT compilation result. ASSERT_OK_AND_ASSIGN(std::string serialized_aot_result, aot_results[0]->SerializeAsString()); ASSERT_OK_AND_ASSIGN( std::unique_ptr aot_result, - compiler->LoadAotCompilationResult(serialized_aot_result)); + compiler()->LoadAotCompilationResult(serialized_aot_result)); // Load Executable from AOT compilation result. ASSERT_OK_AND_ASSIGN( std::unique_ptr executable, std::move(*aot_result) - .LoadExecutable(compiler->PlatformId(), - stream_exec->GetDeviceDescription(), debug_options_)); + .LoadExecutable(stream_executor_platform_id(), device_description(), + debug_options_)); } namespace { @@ -211,11 +196,8 @@ std::string CreateTritonCustomCallBackendConfig() { } // namespace TEST_F(GpuAotCompilationTest, ExportAndLoadExecutableWithTriton) { - auto triton_support = - EnsureTritonSupportsComputeCapability(backend() - .default_stream_executor() - ->GetDeviceDescription() - .gpu_compute_capability()); + auto triton_support = EnsureTritonSupportsComputeCapability( + device_description().gpu_compute_capability()); if (!triton_support.ok()) { GTEST_SKIP() << triton_support; } @@ -237,45 +219,44 @@ TEST_F(GpuAotCompilationTest, ExportAndLoadExecutableWithTriton) { ASSERT_OK_AND_ASSIGN(std::unique_ptr module, ParseAndReturnVerifiedModule(hlo_string)); - auto compiler = backend().compiler(); - auto platform_name = - absl::AsciiStrToUpper(PlatformUtil::CanonicalPlatformName("gpu").value()); - ASSERT_OK_AND_ASSIGN(se::Platform * platform, - se::PlatformManager::PlatformWithName(platform_name)); - ASSERT_OK_AND_ASSIGN(se::StreamExecutor * stream_exec, - platform->ExecutorForDevice(0)); - // Compile AOT. - AotCompilationOptions aot_options(compiler->PlatformId()); - aot_options.set_executor(stream_exec); + AotCompilationOptions aot_options(stream_executor_platform_id()); + aot_options.set_gpu_topology( + GetSingleDeviceGpuTopology("", gpu_target_config())); ASSERT_OK_AND_ASSIGN( std::vector> aot_results, - compiler->CompileAheadOfTime(std::move(module), aot_options)); + compiler()->CompileAheadOfTime(std::move(module), aot_options)); // Serialize-deserialize AOT compilation result. ASSERT_OK_AND_ASSIGN(std::string serialized_aot_result, aot_results[0]->SerializeAsString()); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr aot_result, - compiler->LoadAotCompilationResult(serialized_aot_result)); - // Load Executable from AOT compilation result. + // Load and execute via PjRt test runner. + ExecutableAndOptionsProto proto; + *proto.mutable_serialized_executable() = std::move(serialized_aot_result); + proto.set_pjrt_client_name("PjRtStreamExecutorClient"); + proto.mutable_compile_options() + ->mutable_executable_build_options() + ->set_num_replicas(1); + proto.mutable_compile_options() + ->mutable_executable_build_options() + ->set_num_partitions(1); + std::string serialized_split_proto; + ASSERT_OK(WriteSplitExecutableAndOptions( + proto, + std::make_unique>(&serialized_split_proto))); ASSERT_OK_AND_ASSIGN( - std::unique_ptr executable, - std::move(*aot_result) - .LoadExecutable(compiler->PlatformId(), - stream_exec->GetDeviceDescription(), debug_options_)); - std::unique_ptr wrapped_executable = - test_runner_as_hlo_runner().WrapExecutable(std::move(executable)); + std::unique_ptr executable, + test_runner().DeserializeExecutable(serialized_split_proto)); const xla::Literal literal_1 = xla::LiteralUtil::CreateR0(1.0f); const xla::Literal literal_2 = xla::LiteralUtil::CreateR0(2.0f); const xla::Literal literal_3 = xla::LiteralUtil::CreateR0(3.0f); ASSERT_OK_AND_ASSIGN(Literal result, - test_runner_as_hlo_runner().ExecuteWithExecutable( - wrapped_executable.get(), {&literal_1, &literal_3})); + test_runner().ExecuteWithExecutable( + executable.get(), {&literal_1, &literal_3})); EXPECT_TRUE(LiteralTestUtil::Equal( LiteralUtil::MakeTuple({&literal_2, &literal_3}), result)); diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc index 4ca8fef163aebd..a5408c0f2f3587 100644 --- a/third_party/xla/xla/service/gpu/gpu_compiler.cc +++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc @@ -2379,16 +2379,13 @@ absl::StatusOr> GpuCompiler::RunHloPasses( DumpHloModuleMetadataIfEnabled(module.get()); - AutotuneResults autotune_results; if (stream_exec != nullptr) { - ABSL_RETURN_IF_ERROR( - AutotunerCache::SerializeAutotuneResults(&autotune_results)); ABSL_RETURN_IF_ERROR(SerializeAutotuneResultsToFile(debug_opts)); } std::optional optimized_fingerprint; if (should_upload_hlo_modules) { optimized_fingerprint = - MaybeUploadOptimizedGpuSymbols(module.get(), autotune_results); + MaybeUploadOptimizedGpuSymbols(module.get(), AutotuneResults()); } if (unoptimized_fingerprint.has_value() && optimized_fingerprint.has_value()) { diff --git a/third_party/xla/xla/service/memory_space_assignment/BUILD b/third_party/xla/xla/service/memory_space_assignment/BUILD index f0aa69e553c130..8595149933f3be 100644 --- a/third_party/xla/xla/service/memory_space_assignment/BUILD +++ b/third_party/xla/xla/service/memory_space_assignment/BUILD @@ -105,7 +105,6 @@ xla_test( ":utils", "//xla:comparison_util", "//xla:literal_util", - "//xla:shape_tree", "//xla:shape_util", "//xla:util", "//xla:xla_data_proto_cc", @@ -113,6 +112,7 @@ xla_test( "//xla/hlo/analysis:hlo_dataflow_analysis", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:verified_hlo_module", + "//xla/hlo/transforms/simplifiers:instruction_hoister", "//xla/hlo/utils:hlo_live_range", "//xla/hlo/utils:hlo_matchers", "//xla/service:hlo_buffer", @@ -124,10 +124,8 @@ xla_test( "//xla/service/heap_simulator:allocation_block", "//xla/tests:test_utils", "//xla/tests:xla_internal_test_main", - "//xla/tsl/lib/core:status_test_util", "//xla/tsl/platform:errors", "//xla/tsl/platform:logging", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", @@ -236,7 +234,6 @@ cc_library( "@highwayhash", "@highwayhash//:arch_specific", "@highwayhash//:hh_types", - "@tsl//tsl/platform:statusor", ], ) diff --git a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc index abf40b148bf0f9..f01572da6645f8 100644 --- a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc +++ b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc @@ -2210,9 +2210,13 @@ void MsaAlgorithm::CreateAllocationValuesForJointProcessedValues( defining_instruction->users().end(), may_be_replaced_by_slice_fn); if (!may_be_replaced_by_slice) { - VLOG(3) << "Skip " << interval.buffer->ToShortString() - << " because the buffer is larger than the heap size."; - continue; + if (!MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( + options_.msa_tensor_overrides, + interval.buffer->defining_position(), interval.size)) { + VLOG(3) << "Skip " << interval.buffer->ToShortString() + << " because the buffer is larger than the heap size."; + continue; + } } } @@ -5412,6 +5416,11 @@ void MsaAlgorithm::MaybeSplitAllocationValues( bool MsaAlgorithm::RequiresNoCopyAlternateMemAllocation( AllocationValue& allocation_value) const { + if (MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( + options_.msa_tensor_overrides, allocation_value.defining_position(), + allocation_value.size())) { + return true; + } return allocation_value.value()->shape().has_layout() && allocation_value.value()->shape().layout().memory_space() == options_.alternate_memory_space; @@ -5420,7 +5429,10 @@ bool MsaAlgorithm::RequiresNoCopyAlternateMemAllocation( void MsaAlgorithm::AssignDefaultMemIfNotAllowedInAlternateMem( AllocationValue& allocation_value, int64_t definition_time) { if (!options_.is_position_allowed_in_alternate_mem_fn( - allocation_value.defining_position())) { + allocation_value.defining_position()) || + MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( + options_.msa_tensor_overrides, allocation_value.defining_position(), + allocation_value.size())) { std::optional existing_req = RequiredMemoryAssignmentAt(allocation_value.value(), definition_time); // If a value is pre-colored or already possesses an explicit alternate @@ -5697,6 +5709,15 @@ absl::StatusOr MsaAlgorithm::AllocateAllocationValues( << RequiresNoCopyAlternateMemAllocation(allocation_value); if (RequiresNoCopyAlternateMemAllocation(allocation_value) && allocation_value.size() > available_heap_size()) { + if (MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( + options_.msa_tensor_overrides, + allocation_value.defining_position(), allocation_value.size())) { + return absl::ResourceExhaustedError(absl::StrCat( + "Cannot allocate pinned tensor in alternate memory: tensor size (", + allocation_value.size(), " bytes) exceeds available heap size (", + available_heap_size(), " bytes) for defining instruction '", + allocation_value.defining_instruction()->name(), "'")); + } VLOG(3) << "Skip " << allocation_value.value()->ToShortString() << " because the buffer is larger than the heap size."; continue; @@ -5830,6 +5851,79 @@ absl::StatusOr MsaAlgorithm::AllocateAllocationValues( request, result, options_.prefetch_interval_picker->retry_number()); } + + if (request.fail_on_unsatisfied_override || request.strict_timing || + (request.require_no_copy_alternate_mem_allocation && + MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( + options_.msa_tensor_overrides, + allocation_value_to_update.defining_position(), + allocation_value_to_update.size()))) { + if (allocate_segment_result != AllocationResult::kSuccess) { + std::string reason = ResultToString(allocate_segment_result); + if (result_is(allocate_segment_result, + AllocationResult::kFailOutOfMemory)) { + reason = "Out of alternate memory / spatial placement conflict"; + } else if (result_is( + allocate_segment_result, + AllocationResult::kFailViolatesAsyncCopyResource)) { + reason = "Copy resource limit reached / bandwidth saturation"; + } else if (result_is(allocate_segment_result, + AllocationResult::kFailOutOfAsyncCopies)) { + reason = "Ran out of outstanding asynchronous copies"; + } else if (result_is(allocate_segment_result, + AllocationResult::kFailLiveRangeTooShort)) { + reason = "Live range too short / dependency ordering"; + } + return absl::FailedPreconditionError(absl::StrCat( + "MSA strict override failed for instruction '", + request.use->hlo_use.instruction->name(), "' operand ", + request.use->hlo_use.operand_number, ", shape ", + request.use->hlo_use.instruction + ->operand(request.use->hlo_use.operand_number) + ->shape() + .ToString(), + ", defining instruction '", + allocation_value_to_update.defining_instruction()->name(), + "': ", reason, " (requested schedule time: ", + request.preferred_prefetch_time.has_value() + ? absl::StrCat(*request.preferred_prefetch_time) + : "n/a", + ")")); + } + if (allocate_segment_result == AllocationResult::kSuccess && + request.strict_timing && + request.preferred_prefetch_time.has_value()) { + const Allocation* allocation = + allocation_sequence->empty() + ? nullptr + : allocation_sequence->back().get(); + int64_t actual_prefetch_time = 0; + if (allocation && allocation->is_copy_allocation()) { + actual_prefetch_time = + static_cast(allocation) + ->copy_start_schedule_after(); + } else if (allocation && allocation->is_sliced_copy_allocation()) { + actual_prefetch_time = + static_cast(allocation) + ->slice_details_sorted_by_start_time() + .front() + .copy_start_after_time; + } + if (actual_prefetch_time != *request.preferred_prefetch_time) { + return absl::FailedPreconditionError(absl::StrCat( + "MSA strict timing override unsatisfied for instruction '", + request.use->hlo_use.instruction->name(), "' operand ", + request.use->hlo_use.operand_number, ", shape ", + request.use->hlo_use.instruction + ->operand(request.use->hlo_use.operand_number) + ->shape() + .ToString(), + ": scheduled prefetch time (", actual_prefetch_time, + ") does not match requested prefetch time (", + *request.preferred_prefetch_time, ")")); + } + } + } if (allocate_segment_result == AllocationResult::kSuccess && NeedsMirroredAllocation(allocation_value_to_update, use, previous_use)) { @@ -6256,28 +6350,61 @@ AllocationRequest MsaAlgorithm::CreateAllocationRequest( } } + bool strict_timing = false; + bool fail_on_unsatisfied_override = false; + bool is_prefetch_override = false; + bool require_end_colored_in_default_memory = false; + + if (MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( + options_.msa_tensor_overrides, hlo_use, allocation_value.size()) || + MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( + options_.msa_tensor_overrides, allocation_value.defining_position(), + allocation_value.size())) { + allow_prefetch = false; + allow_no_copy_alternate_mem_allocation = false; + require_end_colored_in_default_memory = true; + AddRequiredAssignment( + allocation_value_to_update.value(), hlo_use.instruction, + MemorySpace::kDefault, use_time, + RequiredMemoryAssignment::Source::kUseNotAllowedInAlternateMemory); + } + if (MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( + options_.msa_tensor_overrides, allocation_value.defining_position(), + allocation_value.size())) { + require_no_copy_alternate_mem_allocation = true; + allow_prefetch = false; + } + int64_t live_range_start_time = (earliest_prefetch_time.has_value() ? earliest_prefetch_time.value() : std::min(definition_time, use_time)); - auto overridden_preferred_prefetch_time = - MemorySpaceAssignmentUtils::GetOverriddenPreferredPrefetchTime( + auto prefetch_override_info = + MemorySpaceAssignmentUtils::GetPrefetchOverrideInfo( + options_.msa_tensor_overrides, options_.preferred_prefetch_overrides, allocation_value.size(), hlo_use, instruction_schedule, live_range_start_time, - latest_prefetch_time); - CHECK_OK(overridden_preferred_prefetch_time.status()); - if (overridden_preferred_prefetch_time.value().has_value()) { - VLOG(1) << "Overriding preferred prefetch for " - << hlo_use.instruction->name() << " operand number " - << hlo_use.operand_number << " operand index " - << hlo_use.operand_index.ToString() << " size " - << allocation_value.size() << " live range (" + latest_prefetch_time, allocation_value.defining_position()); + CHECK_OK(prefetch_override_info.status()); + if (prefetch_override_info.value().has_value()) { + VLOG(1) << "Overriding prefetch for " << hlo_use.instruction->name() + << " operand number " << hlo_use.operand_number + << " operand index " << hlo_use.operand_index.ToString() + << " size " << allocation_value.size() << " live range (" << live_range_start_time << ", " << latest_prefetch_time << ") from " << (preferred_prefetch_time.has_value() ? preferred_prefetch_time.value() : -1) - << " to " << overridden_preferred_prefetch_time.value().value(); - preferred_prefetch_time = overridden_preferred_prefetch_time.value(); + << " to " + << (prefetch_override_info->value().prefetch_time.has_value() + ? absl::StrCat( + *prefetch_override_info->value().prefetch_time) + : "none"); + preferred_prefetch_time = prefetch_override_info->value().prefetch_time; + strict_timing = prefetch_override_info->value().strict_timing; + fail_on_unsatisfied_override = + prefetch_override_info->value().fail_on_unsatisfied_override; + is_prefetch_override = true; } // Rarely, (e.g., when conditional true and false parameters are the @@ -6306,6 +6433,11 @@ AllocationRequest MsaAlgorithm::CreateAllocationRequest( request.required_copy_allocation_for = required_copy_allocation_for; request.required_copy_for_slice = required_copy_for_slice; request.allocation_value_to_update = &allocation_value_to_update; + request.strict_timing = strict_timing; + request.fail_on_unsatisfied_override = fail_on_unsatisfied_override; + request.is_prefetch_override = is_prefetch_override; + request.require_end_colored_in_default_memory = + require_end_colored_in_default_memory; } if (shape_override.has_value()) { @@ -8627,7 +8759,13 @@ AllocationResult MsaAlgorithm::AllocateSegment(AllocationRequest& request) { request.inclusive_start_time); std::optional required_memory_space_at_start; if (required_assignment_at_start.has_value()) { - required_memory_space_at_start = required_assignment_at_start->memory_space; + if (request.require_no_copy_alternate_mem_allocation && + required_assignment_at_start->memory_space == MemorySpace::kDefault) { + required_assignment_at_start = std::nullopt; + } else { + required_memory_space_at_start = + required_assignment_at_start->memory_space; + } } // Find required assignment both for the use and its aliases. If they are both // non-nullopt, then make sure they require the same assignment. @@ -8736,8 +8874,9 @@ AllocationResult MsaAlgorithm::AllocateSegment(AllocationRequest& request) { // First try keeping the allocation entirely in the alternate memory. if (!request.require_start_colored_in_default_memory && !request.require_end_colored_in_default_memory && - required_memory_space_at_start != MemorySpace::kDefault && - required_memory_space_at_end != MemorySpace::kDefault && + (request.require_no_copy_alternate_mem_allocation || + (required_memory_space_at_start != MemorySpace::kDefault && + required_memory_space_at_end != MemorySpace::kDefault)) && request.allow_no_copy_alternate_mem_allocation && !request.require_copy_allocation) { CheckAndUpdateForDualLiveAllocationValues(required_assignment_at_start, @@ -8752,7 +8891,11 @@ AllocationResult MsaAlgorithm::AllocateSegment(AllocationRequest& request) { } } - CHECK(!request.require_no_copy_alternate_mem_allocation); + if (request.require_no_copy_alternate_mem_allocation) { + return allocation_result != AllocationResult::kSuccess + ? allocation_result + : AllocationResult::kFailOutOfMemory; + } if (request.require_start_colored_in_alternate_memory) { // Since no-copy-allocation failed, continuous allocation is not possible in @@ -8916,6 +9059,9 @@ AllocationResult MsaAlgorithm::AllocateSegment(AllocationRequest& request) { << ") doesn't match the preferred prefetch time (" << *request.preferred_prefetch_time << "): " << request.use->hlo_use.ToString(); + if (request.strict_timing) { + return AllocationResult::kFailLiveRangeTooShort; + } } } return AllocationResult::kSuccess; @@ -8928,6 +9074,9 @@ AllocationResult MsaAlgorithm::AllocateSegment(AllocationRequest& request) { << ") which could not be satisfied: " << request.use->hlo_use.ToString(); } + if (request.fail_on_unsatisfied_override || request.strict_timing) { + return prefetch_result; + } result_mark(prefetch_result, allocation_result); } @@ -9177,9 +9326,11 @@ AllocationResult MsaAlgorithm::AllocateInAlternateMemoryNoCopy( bool can_eliminate_copy = false; if (request.allocation_value->allocation_sequence()->empty()) { // There hasn't been any allocations for this interval so far. We can - // eliminate copy if the value can be placed in the alternate memory. - can_eliminate_copy = options_.is_allowed_in_alternate_mem_fn( - *request.allocation_value->value()); + // eliminate copy if the value can be placed in the alternate memory or is + // pinned. + can_eliminate_copy = request.require_no_copy_alternate_mem_allocation || + options_.is_allowed_in_alternate_mem_fn( + *request.allocation_value->value()); } else { // If there has been a previous allocation, we can eliminate the copy if the // previous allocation was also in the alternate memory. @@ -10115,12 +10266,24 @@ AllocationResult MsaAlgorithm::InitializePrefetchIntervalPicker( std::optional preferred_prefetch_time = context.request->preferred_prefetch_time; if (preferred_prefetch_time) { - preferred_prefetch_time = - std::max(*preferred_prefetch_time, earliest_exclusive_prefetch_time); + if (context.request->strict_timing) { + if (*preferred_prefetch_time < earliest_exclusive_prefetch_time || + *preferred_prefetch_time >= context.prefetch_end_time) { + VLOG(3) << "Strict prefetch time " << *preferred_prefetch_time + << " is outside allowable prefetch range (" + << earliest_exclusive_prefetch_time << ", " + << context.prefetch_end_time << ")."; + return AllocationResult::kFailLiveRangeTooShort; + } + } else { + preferred_prefetch_time = + std::max(*preferred_prefetch_time, earliest_exclusive_prefetch_time); + } } options_.prefetch_interval_picker->Begin( context.request->use->hlo_use, earliest_exclusive_prefetch_time, - context.prefetch_end_time, preferred_prefetch_time); + context.prefetch_end_time, preferred_prefetch_time, + context.request->strict_timing); VLOG(3) << "Trying prefetch picker = " << options_.prefetch_interval_picker->ToDebugString(); diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation_value.cc b/third_party/xla/xla/service/memory_space_assignment/allocation_value.cc index bb372b08055c6e..d1df3be48dccde 100644 --- a/third_party/xla/xla/service/memory_space_assignment/allocation_value.cc +++ b/third_party/xla/xla/service/memory_space_assignment/allocation_value.cc @@ -94,14 +94,18 @@ std::string AllocationRequest::ToString() const { no_copy_chunk_inclusive_start_time.has_value() ? absl::StrCat(*no_copy_chunk_inclusive_start_time) : "nullopt"), - absl::StrCat("require_start_colored_in_alternate_memmory: ", - require_start_colored_in_alternate_memory, - "; require_end_colored_in_alternate_memory: ", - require_end_colored_in_alternate_memory, - "; require_start_colored_in_default_memory: ", - require_start_colored_in_default_memory, - "; require_end_colored_in_default_memory: ", - require_end_colored_in_default_memory)}, + absl::StrCat( + "require_start_colored_in_alternate_memmory: ", + require_start_colored_in_alternate_memory, + "; require_end_colored_in_alternate_memory: ", + require_end_colored_in_alternate_memory, + "; require_start_colored_in_default_memory: ", + require_start_colored_in_default_memory, + "; require_end_colored_in_default_memory: ", + require_end_colored_in_default_memory, + "; strict_timing: ", strict_timing, + "; fail_on_unsatisfied_override: ", fail_on_unsatisfied_override, + "; is_prefetch_override: ", is_prefetch_override)}, "\n"); } diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation_value.h b/third_party/xla/xla/service/memory_space_assignment/allocation_value.h index ce623e08eb4c01..502335244f6b37 100644 --- a/third_party/xla/xla/service/memory_space_assignment/allocation_value.h +++ b/third_party/xla/xla/service/memory_space_assignment/allocation_value.h @@ -291,6 +291,13 @@ struct AllocationRequest { // Indicates if the AllocationRequest end time (use time) has a default // memory color requirement. bool require_end_colored_in_default_memory = false; + // If true, forces the prefetch to start at the exact preferred_prefetch_time + // without searching outwards. + bool strict_timing = false; + // If true, any failure to satisfy this override results in a fatal error. + bool fail_on_unsatisfied_override = false; + // Indicates if this request has a prefetch override. + bool is_prefetch_override = false; std::string ToString() const; }; diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.proto b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.proto index 0f1ed29b50357a..db0f1949eacfe9 100644 --- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.proto +++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.proto @@ -113,7 +113,13 @@ message PreferredPrefetchOverrideOptions { // Preferred prefetch time is set to before the first instruction that // matches the filter. HloPositionMatcher before_instruction = 5; + // Preferred prefetch time set directly to a logical schedule time. + int64 logical_time = 6; } + // If true, the picker and allocator must ONLY attempt this exact time (no + // outward search). If it cannot be scheduled at this exact time, fail + // compilation. + optional bool strict_timing = 7; } // Filters operands in an HLO schedule and overrides preferred prefetch times @@ -123,12 +129,16 @@ message PreferredPrefetchOverride { optional HloOperandFilter hlo_operand_filter = 1; optional xla.memory_space_assignment.PreferredPrefetchOverrideOptions override_options = 2; + // If true, any failure to fulfill the action results in a fatal error. + optional bool fail_on_unsatisfied_override = 3; } // Encloses chained override configs. The first config has highest precedence // and so on. message PreferredPrefetchOverrides { repeated PreferredPrefetchOverride overrides = 1; + // Global flag to fail compilation if any override is unsatisfied. + optional bool fail_on_unsatisfied_override = 2; } // Specifies details on how to randomally select HloPositions for perturbation. @@ -213,6 +223,32 @@ message MsaSortOrderOverrides { repeated MsaSortOrderOverride overrides = 1; } +// Action to pin a buffer in alternate memory across its entire lifetime. +message PinInAlternateMemoryAction {} + +// Action to keep a buffer or operand use strictly in default memory (HBM). +message KeepInDefaultMemoryAction {} + +// Unified override rule for tensor and operand placement and timing. Overrides +// specified via MsaTensorOverride are always strictly enforced; compilation +// fails if an override cannot be fulfilled. +message MsaTensorOverride { + // Target filter for operand use or buffer position. + optional HloOperandFilter hlo_operand_filter = 1; + optional HloPositionMatcher hlo_position_matcher = 2; + + oneof action { + PreferredPrefetchOverrideOptions prefetch = 3; + PinInAlternateMemoryAction pin_in_alternate_memory = 4; + KeepInDefaultMemoryAction keep_in_default_memory = 5; + } +} + +// Encloses chained tensor override configs. +message MsaTensorOverrides { + repeated MsaTensorOverride overrides = 1; +} + // Expanded scoped alternate memory is a feature used at the end of MSA, in // in which we attempt to expand the size of allocated scoped alternate memory // buffers to the largest contiguous open space available. diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc index 5664b984154742..cdd421527b137a 100644 --- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc +++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc @@ -55,6 +55,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_print_options.h" #include "xla/hlo/ir/hlo_schedule.h" #include "xla/hlo/testlib/verified_hlo_module.h" +#include "xla/hlo/transforms/simplifiers/instruction_hoister.h" #include "xla/hlo/utils/hlo_live_range.h" #include "xla/hlo/utils/hlo_matchers.h" #include "xla/layout.h" @@ -82,13 +83,10 @@ limitations under the License. #include "xla/service/memory_space_assignment/testing_utils.h" #include "xla/service/memory_space_assignment/utils.h" #include "xla/shape.h" -#include "xla/shape_tree.h" #include "xla/shape_util.h" #include "xla/tests/test_utils.h" -#include "xla/tsl/lib/core/status_test_util.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/logging.h" -#include "xla/tsl/platform/statusor.h" #include "xla/util.h" #include "xla/xla_data.pb.h" #include "tsl/platform/protobuf.h" // IWYU pragma: keep @@ -2436,6 +2434,529 @@ TEST_F(MemorySpaceAssignmentTest, FilterUpdatePreferredPrefetchNoMatchTest) { EXPECT_THAT(sequence.instructions()[10], op::CopyDone()); } +TEST_F(MemorySpaceAssignmentTest, StrictPrefetchExactTimeSuccess) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* negate2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); + HloInstruction* negate3 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); + HloInstruction* negate4 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); + HloInstruction* negate5 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); + HloInstruction* negate6 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, + negate3, negate4, negate5, negate6, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + + const std::string text_proto = R"pb( + overrides { + hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } + prefetch { + after_instruction { instruction_name_regex: ".*negate.3.*" } + strict_timing: true + } + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpace(module.get(), std::move(options)); + + EXPECT_THAT(add, op::Add(op::Negate(), op::AsyncCopy(kAlternateMemorySpace, + kDefaultMemorySpace, + op::Parameter(1)))); + // Ensure the CopyStart is scheduled after negate.3 (schedule index 6). + const HloInstructionSequence& sequence = + module->schedule().sequence(computation); + EXPECT_THAT(sequence.instructions()[0], op::Parameter(0)); + EXPECT_THAT(sequence.instructions()[1], op::Parameter(1)); + EXPECT_THAT(sequence.instructions()[5], op::Negate()); + EXPECT_THAT(sequence.instructions()[6], op::CopyStart()); + EXPECT_THAT(sequence.instructions()[10], op::CopyDone()); +} + +TEST_F(MemorySpaceAssignmentTest, StrictPrefetchLogicalTimeSuccess) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* negate2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); + HloInstruction* negate3 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); + HloInstruction* negate4 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); + HloInstruction* negate5 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); + HloInstruction* negate6 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, + negate3, negate4, negate5, negate6, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + + const std::string text_proto = R"pb( + overrides { + hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } + prefetch { logical_time: 4 strict_timing: true } + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpace(module.get(), std::move(options)); + + const HloInstructionSequence& sequence = + module->schedule().sequence(computation); + EXPECT_THAT(sequence.instructions()[5], op::CopyStart()); +} + +TEST_F(MemorySpaceAssignmentTest, StrictPrefetchUnsatisfiableFails) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate1, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + const std::string text_proto = R"pb( + overrides { + hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } + prefetch { logical_time: 0 strict_timing: true } + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + InstructionHoister instruction_hoister; + CHECK_OK(instruction_hoister.Run(module.get()).status()); + InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 10); + auto status_or = AssignMemorySpaceAndReturnStatus( + module.get(), std::move(options), /*buffer_interval_compare=*/{}, + &prefetch_interval_picker); + EXPECT_FALSE(status_or.ok()); +} + +TEST_F(MemorySpaceAssignmentTest, StrictPinInAlternateMemory) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate1, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + const std::string text_proto = R"pb( + overrides { + hlo_position_matcher { instruction_name_regex: ".*p1.*" } + pin_in_alternate_memory {} + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpace(module.get(), std::move(options)); + + // p1 is pinned in alternate memory, so no copy allocation is created and + // add consumes p1 directly. + EXPECT_THAT(add, op::Add(op::Negate(), op::Parameter(1))); + EXPECT_EQ(p1->shape().layout().memory_space(), kAlternateMemorySpace); +} + +TEST_F(MemorySpaceAssignmentTest, StrictKeepInDefaultMemory) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* negate2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); + HloInstruction* negate3 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); + HloInstruction* negate4 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); + HloInstruction* negate5 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); + HloInstruction* negate6 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, + negate3, negate4, negate5, negate6, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + // Strictly keep p1 in default memory, preventing prefetch. + const std::string text_proto = R"pb( + overrides { + hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } + keep_in_default_memory {} + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpace(module.get(), std::move(options)); + + // Without override, p1 would be prefetched via AsyncCopy. + // With keep_in_default_memory, add consumes Parameter(1) directly. + EXPECT_THAT(add, op::Add(op::Negate(), op::Parameter(1))); + EXPECT_EQ(p1->shape().layout().memory_space(), kDefaultMemorySpace); +} + +TEST_F(MemorySpaceAssignmentTest, StrictKeepInDefaultMemoryDefiningPosition) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* negate2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); + HloInstruction* negate3 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); + HloInstruction* negate4 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); + HloInstruction* negate5 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); + HloInstruction* negate6 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, + negate3, negate4, negate5, negate6, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + const std::string text_proto = R"pb( + overrides { + hlo_position_matcher { instruction_name_regex: ".*p1.*" } + keep_in_default_memory {} + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpace(module.get(), std::move(options)); + + EXPECT_THAT(add, op::Add(op::Negate(), op::Parameter(1))); + EXPECT_EQ(p1->shape().layout().memory_space(), kDefaultMemorySpace); +} + +TEST_F(MemorySpaceAssignmentTest, + PositionMatcherDoesNotMatchConsumerInstruction) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* negate2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); + HloInstruction* negate3 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); + HloInstruction* negate4 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); + HloInstruction* negate5 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); + HloInstruction* negate6 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, + negate3, negate4, negate5, negate6, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + // Matching ".*add.*" with HloPositionMatcher should ONLY match buffers + // PRODUCED by add, NOT buffers CONSUMED by add (like p1). + const std::string text_proto = R"pb( + overrides { + hlo_position_matcher { instruction_name_regex: ".*add.*" } + keep_in_default_memory {} + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpace(module.get(), std::move(options)); + + // Since p1's defining position is p1 (not add), it is NOT kept in default + // memory and gets prefetched into alternate memory space. + EXPECT_THAT(add, op::Add(op::Negate(), op::AsyncCopy(kAlternateMemorySpace, + kDefaultMemorySpace, + op::Parameter(1)))); +} + +TEST_F(MemorySpaceAssignmentTest, StrictPrefetchBeforeInstructionSuccess) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* negate2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); + HloInstruction* negate3 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); + HloInstruction* negate4 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); + HloInstruction* negate5 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); + HloInstruction* negate6 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, + negate3, negate4, negate5, negate6, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + + const std::string text_proto = R"pb( + overrides { + hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } + prefetch { + before_instruction { instruction_name_regex: ".*negate.4.*" } + strict_timing: true + } + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpace(module.get(), std::move(options)); + + const HloInstructionSequence& sequence = + module->schedule().sequence(computation); + EXPECT_THAT(sequence.instructions()[6], op::CopyStart()); +} + +TEST_F(MemorySpaceAssignmentTest, StrictPrefetchExactTimeWithStrictBool) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* negate2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); + HloInstruction* negate3 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); + HloInstruction* negate4 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); + HloInstruction* negate5 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); + HloInstruction* negate6 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, + negate3, negate4, negate5, negate6, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + + const std::string text_proto = R"pb( + overrides { + hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } + prefetch { logical_time: 4 strict_timing: true } + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpace(module.get(), std::move(options)); + + const HloInstructionSequence& sequence = + module->schedule().sequence(computation); + EXPECT_THAT(sequence.instructions()[5], op::CopyStart()); +} + +TEST_F(MemorySpaceAssignmentTest, StrictPrefetchCostAnalysisExactTimeSuccess) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* negate2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); + HloInstruction* negate3 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); + HloInstruction* negate4 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); + HloInstruction* negate5 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); + HloInstruction* negate6 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, + negate3, negate4, negate5, negate6, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + + const std::string text_proto = R"pb( + overrides { + hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } + prefetch { + after_instruction { instruction_name_regex: ".*negate.3.*" } + strict_timing: true + } + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpaceUsingCostAnalysis(module.get(), std::move(options)); + + const HloInstructionSequence& sequence = + module->schedule().sequence(computation); + EXPECT_THAT(sequence.instructions()[6], op::CopyStart()); +} + +TEST_F(MemorySpaceAssignmentTest, StrictPinInAlternateMemoryFailsOnOom) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); // 24 bytes + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, negate0}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + options.max_size_in_bytes = 10; + const std::string text_proto = R"pb( + overrides { + hlo_position_matcher { instruction_name_regex: ".*p0.*" } + pin_in_alternate_memory {} + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + InstructionHoister instruction_hoister; + CHECK_OK(instruction_hoister.Run(module.get()).status()); + InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 10); + auto status_or = AssignMemorySpaceAndReturnStatus( + module.get(), std::move(options), /*buffer_interval_compare=*/{}, + &prefetch_interval_picker); + EXPECT_FALSE(status_or.ok()); +} + TEST_F(MemorySpaceAssignmentTest, EvictAndPrefetch) { std::unique_ptr module = CreateEvictAndPrefetchModule(); diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test_base.h b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test_base.h index 8f869d9345e346..936607a2d2dc0b 100644 --- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test_base.h +++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test_base.h @@ -304,6 +304,12 @@ class MemorySpaceAssignmentTestBase : public HloTestBase { if (options_override) { options = *std::move(options_override); } + for (const auto& override : options.msa_tensor_overrides.overrides()) { + if (override.has_pin_in_alternate_memory()) { + check_parameters_in_default_memory = false; + break; + } + } std::unique_ptr test_comparator; if (buffer_interval_compare.has_value()) { test_comparator = std::make_unique( diff --git a/third_party/xla/xla/service/memory_space_assignment/options.cc b/third_party/xla/xla/service/memory_space_assignment/options.cc index 8bbab64a6878ad..5bc554c0574cef 100644 --- a/third_party/xla/xla/service/memory_space_assignment/options.cc +++ b/third_party/xla/xla/service/memory_space_assignment/options.cc @@ -105,6 +105,8 @@ std::string Options::ToString() const { autotuning_config.has_value() ? "present" : "nullopt"), absl::StrCat("preferred_prefetch_overrides: \n", preferred_prefetch_overrides.DebugString()), + absl::StrCat("msa_tensor_overrides: \n", + msa_tensor_overrides.DebugString()), absl::StrCat("sliced_prefetch_options: \n", sliced_prefetch_options.DebugString()), absl::StrCat("memory_bound_loop_optimizer_options: \n", diff --git a/third_party/xla/xla/service/memory_space_assignment/options.h b/third_party/xla/xla/service/memory_space_assignment/options.h index b2635a1019bcc8..47fb95cc12ef17 100644 --- a/third_party/xla/xla/service/memory_space_assignment/options.h +++ b/third_party/xla/xla/service/memory_space_assignment/options.h @@ -407,6 +407,10 @@ struct Options { // filtered prefetches. PreferredPrefetchOverrides preferred_prefetch_overrides; + // Unified overrides for tensor placement, pinning, retention, and prefetch + // timing. + MsaTensorOverrides msa_tensor_overrides; + // Options for slicing prefetches into smaller asynchronously copied pieces. SlicedPrefetchOptions sliced_prefetch_options; diff --git a/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.cc b/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.cc index d0594b1e9eace1..4a38f949752eed 100644 --- a/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.cc +++ b/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.cc @@ -93,15 +93,22 @@ float InstructionCountPrefetchIntervalPicker::GetLogicalIntervalElapsed( void InstructionCountPrefetchIntervalPicker::Begin( const HloUse& use, int64_t start_time, int64_t end_time, - std::optional preferred_time) { + std::optional preferred_time, bool strict_timing) { end_time_ = end_time; + strict_timing_ = strict_timing; const Shape& shape = ShapeUtil::GetSubshape( use.instruction->operand(use.operand_number)->shape(), use.operand_index); if (preferred_time) { current_prefetch_time_ = *preferred_time; + if (strict_timing) { + strict_prefetch_time_ = *preferred_time; + } } else { current_prefetch_time_ = PreferredPrefetchStartTime(shape, start_time, end_time, end_time); + if (strict_timing) { + strict_prefetch_time_ = current_prefetch_time_; + } } } @@ -112,10 +119,16 @@ int64_t InstructionCountPrefetchIntervalPicker::Next() { } bool InstructionCountPrefetchIntervalPicker::Done() const { + if (strict_timing_) { + return current_prefetch_time_ > strict_prefetch_time_; + } return end_time_ - current_prefetch_time_ <= min_overlap_count_; } int64_t InstructionCountPrefetchIntervalPicker::latest_time() const { + if (strict_timing_) { + return strict_prefetch_time_; + } return end_time_ - min_overlap_count_ - 1; } @@ -356,7 +369,7 @@ int64_t CostAnalysisPrefetchIntervalPicker::EstimatedPrefetchEndTime( void CostAnalysisPrefetchIntervalPicker::Begin( const HloUse& use, int64_t start_time, int64_t end_time, - std::optional preferred_time) { + std::optional preferred_time, bool strict_timing) { const Shape& shape = ShapeUtil::GetSubshape( use.instruction->operand(use.operand_number)->shape(), use.operand_index); int64_t shape_size = size_override_ ? *size_override_ @@ -374,6 +387,17 @@ void CostAnalysisPrefetchIntervalPicker::Begin( end_logical_time_ = end_time; int end_nest_level = computation_nest_level_[end_logical_time_]; + if (strict_timing && preferred_time.has_value()) { + int64_t target_time = *preferred_time; + earliest_prefetch_time_ = target_time; + latest_prefetch_time_ = target_time; + increasing_prefetch_time_iterator_ = target_time; + decreasing_prefetch_time_iterator_ = target_time; + using_increasing_prefetch_time_iterator_ = true; + Next(); + return; + } + // Find the latest time we're allowed to start prefetching. float min_interval = min_overlap_to_async_copy_ratio_ * async_copy_elapsed_; latest_prefetch_time_ = @@ -527,8 +551,11 @@ std::string CostAnalysisPrefetchIntervalPicker::ToDebugString() const { int current_logical_prefetch_time = using_increasing_prefetch_time_iterator_ ? increasing_prefetch_time_iterator_ : decreasing_prefetch_time_iterator_; - float logical_interval_elapsed = GetLogicalIntervalElapsed( - current_logical_prefetch_time, end_logical_time_); + float logical_interval_elapsed = 0.0f; + if (current_logical_prefetch_time <= end_logical_time_) { + logical_interval_elapsed = GetLogicalIntervalElapsed( + current_logical_prefetch_time, end_logical_time_); + } return absl::StrCat( "Async copy elapsed (s) = ", async_copy_elapsed_, ", inst elapsed reduction (s) = ", inst_elapsed_reduction_, diff --git a/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.h b/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.h index e12ab1a729da89..56eb842d1e8ebe 100644 --- a/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.h +++ b/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.h @@ -83,7 +83,8 @@ class PrefetchIntervalPicker { // Begins the iterator for the first start time of the prefetch. virtual void Begin(const HloUse& use, int64_t start_time, int64_t end_time, - std::optional preferred_time) = 0; + std::optional preferred_time, + bool strict_timing = false) = 0; // Advances the start time of the prefetch and returns that value. virtual int64_t Next() = 0; @@ -165,7 +166,8 @@ class InstructionCountPrefetchIntervalPicker : public PrefetchIntervalPicker { int64_t end_time) const override; void Begin(const HloUse& use, int64_t start_time, int64_t end_time, - std::optional preferred_time) override; + std::optional preferred_time, + bool strict_timing = false) override; int64_t Next() override; bool Done() const override; @@ -181,6 +183,8 @@ class InstructionCountPrefetchIntervalPicker : public PrefetchIntervalPicker { int64_t max_overlap_count_; int64_t end_time_; int64_t current_prefetch_time_; + bool strict_timing_ = false; + int64_t strict_prefetch_time_ = -1; }; // Prefetch interval picker that uses cost analysis to overlap asynchronous @@ -231,7 +235,8 @@ class CostAnalysisPrefetchIntervalPicker : public PrefetchIntervalPicker { int64_t end_time) const override; void Begin(const HloUse& use, int64_t start_time, int64_t end_time, - std::optional preferred_time) override; + std::optional preferred_time, + bool strict_timing = false) override; int64_t Next() override; bool Done() const override; diff --git a/third_party/xla/xla/service/memory_space_assignment/utils.cc b/third_party/xla/xla/service/memory_space_assignment/utils.cc index f6b932a08e4b5c..363c3ac4f03364 100644 --- a/third_party/xla/xla/service/memory_space_assignment/utils.cc +++ b/third_party/xla/xla/service/memory_space_assignment/utils.cc @@ -44,7 +44,6 @@ limitations under the License. #include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" #include "xla/shape_util.h" #include "xla/util.h" -#include "tsl/platform/statusor.h" namespace xla { namespace memory_space_assignment { @@ -138,7 +137,7 @@ bool MemorySpaceAssignmentUtils::IsIntervalAllowedInAlternateMemory( }); } -bool MemorySpaceAssignmentUtils::DoesUseMatchFilter( +bool MemorySpaceAssignmentUtils::DoesUseMatchOperandFilter( const HloOperandFilter& filter, const HloUse& hlo_use, int64_t operand_size) { // The order of checks is such that the most expensive checks are done last. @@ -264,6 +263,49 @@ bool MemorySpaceAssignmentUtils::DoesPositionMatchFilter( DoesInstructionMatchRandomFilter(filter, *instruction); } +bool MemorySpaceAssignmentUtils::DoesPositionMatchPositionFilter( + const HloPositionMatcher& filter, const HloPosition& position, + int64_t size) { + if (filter.has_size_gte() && filter.size_gte() > size) { + return false; + } + if (filter.has_size_lte() && filter.size_lte() < size) { + return false; + } + if (filter.has_tuple_index() && + position.index != ShapeIndex(filter.tuple_index().index().begin(), + filter.tuple_index().index().end())) { + return false; + } + return DoesInstructionMatchFilter(filter, *position.instruction) && + DoesInstructionMatchRandomFilter(filter, *position.instruction); +} + +bool MemorySpaceAssignmentUtils::DoesPositionMatchOperandFilter( + const HloOperandFilter& filter, const HloPosition& position, int64_t size) { + if (filter.has_size_gte() && filter.size_gte() > size) { + return false; + } + if (filter.has_size_lte() && filter.size_lte() < size) { + return false; + } + if (filter.has_tuple_index() && + position.index != ShapeIndex(filter.tuple_index().index().begin(), + filter.tuple_index().index().end())) { + return false; + } + if (filter.has_instruction_name_regex() && + !RE2::FullMatch(position.instruction->name(), + filter.instruction_name_regex())) { + return false; + } + if (filter.has_instruction_regex() && + !RE2::FullMatch(position.instruction->ToString(), + filter.instruction_regex())) { + return false; + } + return true; +} bool MemorySpaceAssignmentUtils::DoesInstructionMatchFilter( const HloPositionMatcher& filter, const HloInstruction& instruction) { if (filter.has_instruction_name_regex() && @@ -284,8 +326,8 @@ bool MemorySpaceAssignmentUtils::DoesBufferIntervalMatchHloUseFilter( return true; } for (const HloUse& use : buffer_interval.buffer->GetUses()) { - if (DoesUseMatchFilter(filter.hlo_use_filter(), use, - buffer_interval.size)) { + if (DoesUseMatchOperandFilter(filter.hlo_use_filter(), use, + buffer_interval.size)) { return true; } } @@ -358,6 +400,9 @@ MemorySpaceAssignmentUtils::GetPrefetchTime( case PreferredPrefetchOverrideOptions::kBeforeInstruction: return GetPrefetchTimeBeforeInstruction( override_options.before_instruction(), instruction_schedule); + case PreferredPrefetchOverrideOptions::kLogicalTime: + return static_cast>( + override_options.logical_time()); case PreferredPrefetchOverrideOptions::OPTIONS_NOT_SET: break; } @@ -372,7 +417,7 @@ MemorySpaceAssignmentUtils::GetOverriddenPreferredPrefetchTime( instruction_schedule, int64_t earliest_prefetch_time, int64_t latest_prefetch_time) { for (const auto& override : preferred_prefetch_overrides.overrides()) { - if (!MemorySpaceAssignmentUtils::DoesUseMatchFilter( + if (!MemorySpaceAssignmentUtils::DoesUseMatchOperandFilter( override.hlo_operand_filter(), hlo_use, operand_size)) { continue; } @@ -394,6 +439,130 @@ MemorySpaceAssignmentUtils::GetOverriddenPreferredPrefetchTime( return static_cast>>(std::nullopt); } +absl::StatusOr> +MemorySpaceAssignmentUtils::GetPrefetchOverrideInfo( + const MsaTensorOverrides& msa_tensor_overrides, + const PreferredPrefetchOverrides& preferred_prefetch_overrides, + int64_t operand_size, const HloUse& hlo_use, + const absl::flat_hash_map& + instruction_schedule, + int64_t earliest_prefetch_time, int64_t latest_prefetch_time, + const std::optional& position) { + // First check msa_tensor_overrides. + for (const auto& override : msa_tensor_overrides.overrides()) { + if (!override.has_prefetch()) { + continue; + } + bool matches = false; + if (override.has_hlo_operand_filter() && + DoesUseMatchOperandFilter(override.hlo_operand_filter(), hlo_use, + operand_size)) { + matches = true; + } else if (override.has_hlo_position_matcher() && position.has_value() && + DoesPositionMatchPositionFilter(override.hlo_position_matcher(), + *position, operand_size)) { + matches = true; + } + if (!matches) { + continue; + } + const auto& prefetch_options = override.prefetch(); + ABSL_ASSIGN_OR_RETURN( + auto prefetch_time, + GetPrefetchTime(prefetch_options, earliest_prefetch_time, + latest_prefetch_time, instruction_schedule)); + bool strict_timing = prefetch_options.has_strict_timing() && + prefetch_options.strict_timing(); + PrefetchOverrideInfo info; + info.prefetch_time = prefetch_time; + info.strict_timing = strict_timing; + info.fail_on_unsatisfied_override = true; + return info; + } + + // Next check preferred_prefetch_overrides. + for (const auto& override : preferred_prefetch_overrides.overrides()) { + if (!DoesUseMatchOperandFilter(override.hlo_operand_filter(), hlo_use, + operand_size)) { + continue; + } + ABSL_ASSIGN_OR_RETURN( + auto prefetch_time, + GetPrefetchTime(override.override_options(), earliest_prefetch_time, + latest_prefetch_time, instruction_schedule)); + bool strict_timing = override.override_options().has_strict_timing() && + override.override_options().strict_timing(); + bool fail_on_unsatisfied = + override.fail_on_unsatisfied_override() || + preferred_prefetch_overrides.fail_on_unsatisfied_override() || + strict_timing; + PrefetchOverrideInfo info; + info.prefetch_time = prefetch_time; + info.strict_timing = strict_timing; + info.fail_on_unsatisfied_override = fail_on_unsatisfied; + return info; + } + return std::nullopt; +} + +bool MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( + const MsaTensorOverrides& msa_tensor_overrides, const HloPosition& position, + int64_t size) { + for (const auto& override : msa_tensor_overrides.overrides()) { + if (!override.has_pin_in_alternate_memory()) { + continue; + } + if (override.has_hlo_position_matcher() && + DoesPositionMatchPositionFilter(override.hlo_position_matcher(), + position, size)) { + return true; + } + if (override.has_hlo_operand_filter() && + DoesPositionMatchOperandFilter(override.hlo_operand_filter(), position, + size)) { + return true; + } + } + return false; +} + +bool MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( + const MsaTensorOverrides& msa_tensor_overrides, const HloPosition& position, + int64_t size) { + for (const auto& override : msa_tensor_overrides.overrides()) { + if (!override.has_keep_in_default_memory()) { + continue; + } + if (override.has_hlo_position_matcher() && + DoesPositionMatchPositionFilter(override.hlo_position_matcher(), + position, size)) { + return true; + } + if (override.has_hlo_operand_filter() && + DoesPositionMatchOperandFilter(override.hlo_operand_filter(), position, + size)) { + return true; + } + } + return false; +} + +bool MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( + const MsaTensorOverrides& msa_tensor_overrides, const HloUse& hlo_use, + int64_t operand_size) { + for (const auto& override : msa_tensor_overrides.overrides()) { + if (!override.has_keep_in_default_memory()) { + continue; + } + if (override.has_hlo_operand_filter() && + DoesUseMatchOperandFilter(override.hlo_operand_filter(), hlo_use, + operand_size)) { + return true; + } + } + return false; +} + bool MemorySpaceAssignmentUtils::DoesCrossProgramPrefetchBufferMatchAnyFilter( const MsaSortOrderOverrides& sort_order_overrides, const MsaBufferInterval& buffer_interval) { diff --git a/third_party/xla/xla/service/memory_space_assignment/utils.h b/third_party/xla/xla/service/memory_space_assignment/utils.h index a7bd6f4688dbb3..b8a897252f3e35 100644 --- a/third_party/xla/xla/service/memory_space_assignment/utils.h +++ b/third_party/xla/xla/service/memory_space_assignment/utils.h @@ -18,7 +18,6 @@ limitations under the License. #include #include -#include #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" @@ -36,6 +35,13 @@ namespace memory_space_assignment { using MsaBufferInterval = GlobalDecreasingSizeBestFitHeap::BufferInterval; +// Info resolved from prefetch overrides. +struct PrefetchOverrideInfo { + std::optional prefetch_time; + bool strict_timing = false; + bool fail_on_unsatisfied_override = false; +}; + // Encapsulates common utility methods for memory space assignment. class MemorySpaceAssignmentUtils { public: @@ -55,8 +61,9 @@ class MemorySpaceAssignmentUtils { const HloInstruction& instruction, const absl::flat_hash_set& execution_threads); - static bool DoesUseMatchFilter(const HloOperandFilter& filter, - const HloUse& hlo_use, int64_t operand_size); + static bool DoesUseMatchOperandFilter(const HloOperandFilter& filter, + const HloUse& hlo_use, + int64_t operand_size); static bool DoesInstructionMatchFilter(const HloPositionMatcher& filter, const HloInstruction& instruction); @@ -67,6 +74,14 @@ class MemorySpaceAssignmentUtils { static bool DoesPositionMatchFilter(const HloPositionMatcher& filter, const MsaBufferInterval& buffer_interval); + static bool DoesPositionMatchPositionFilter(const HloPositionMatcher& filter, + const HloPosition& position, + int64_t size); + + static bool DoesPositionMatchOperandFilter(const HloOperandFilter& filter, + const HloPosition& position, + int64_t size); + static absl::StatusOr GetScheduleTimeFromInstructionMatcher( const HloPositionMatcher& position_matcher, @@ -104,6 +119,29 @@ class MemorySpaceAssignmentUtils { instruction_schedule, int64_t earliest_prefetch_time, int64_t latest_prefetch_time); + static absl::StatusOr> + GetPrefetchOverrideInfo( + const MsaTensorOverrides& msa_tensor_overrides, + const PreferredPrefetchOverrides& preferred_prefetch_overrides, + int64_t operand_size, const HloUse& hlo_use, + const absl::flat_hash_map& + instruction_schedule, + int64_t earliest_prefetch_time, int64_t latest_prefetch_time, + const std::optional& position = std::nullopt); + + static bool ShouldPinInAlternateMemory( + const MsaTensorOverrides& msa_tensor_overrides, + const HloPosition& position, int64_t size); + + static bool ShouldKeepInDefaultMemory( + const MsaTensorOverrides& msa_tensor_overrides, + const HloPosition& position, int64_t size); + + static bool ShouldKeepInDefaultMemory( + const MsaTensorOverrides& msa_tensor_overrides, const HloUse& hlo_use, + int64_t operand_size); + static bool DoesCrossProgramPrefetchBufferMatchAnyFilter( const MsaSortOrderOverrides& sort_order_overrides, const MsaBufferInterval& buffer_interval); diff --git a/third_party/xla/xla/stream_executor/sycl/BUILD b/third_party/xla/xla/stream_executor/sycl/BUILD index 4c22875e53d4ad..0511a4e3894db6 100644 --- a/third_party/xla/xla/stream_executor/sycl/BUILD +++ b/third_party/xla/xla/stream_executor/sycl/BUILD @@ -456,6 +456,7 @@ sycl_library( "//xla/tsl/platform:statusor", "@com_google_absl//absl/base", "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/base:no_destructor", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/strings", diff --git a/third_party/xla/xla/stream_executor/sycl/sycl_context.cc b/third_party/xla/xla/stream_executor/sycl/sycl_context.cc index aed9bd0152bb73..6210396b386483 100644 --- a/third_party/xla/xla/stream_executor/sycl/sycl_context.cc +++ b/third_party/xla/xla/stream_executor/sycl/sycl_context.cc @@ -21,7 +21,7 @@ namespace stream_executor::sycl { absl::StatusOr> SyclContext::Create( int device_ordinal) { - ABSL_ASSIGN_OR_RETURN(::sycl::context sycl_context, + ABSL_ASSIGN_OR_RETURN(const ::sycl::context& sycl_context, SyclDevicePool::GetDeviceContext()); return std::make_unique(sycl_context, device_ordinal); } diff --git a/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime.cc b/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime.cc index 81743e08de4a8b..b10b9e6a852d5b 100644 --- a/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime.cc +++ b/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime.cc @@ -186,7 +186,7 @@ bool SyclIsHostMemoryRegistered(const ::sycl::device& device, } } -DevicePool SyclDevicePool::device_pool_; +absl::NoDestructor SyclDevicePool::device_pool_; absl::Status SyclDevicePool::InitDevicePool() { static absl::once_flag device_init_flag; @@ -217,29 +217,32 @@ absl::Status SyclDevicePool::InitDevicePool() { "backend. Check oneAPI installation and environment variables."); return; } - device_pool_ = std::move(devices); + // absl::NoDestructor default-constructs an empty DevicePool, so this + // dereference is safe. + *device_pool_ = std::move(devices); }); return init_status; } -absl::StatusOr<::sycl::context> SyclDevicePool::GetDeviceContext() { +absl::StatusOr SyclDevicePool::GetDeviceContext() { ABSL_RETURN_IF_ERROR(SyclDevicePool::InitDevicePool()); - static ::sycl::context device_context(device_pool_); - return device_context; + // Leaked for the same reason as SyclDevicePool::device_pool_. + static absl::NoDestructor<::sycl::context> device_context(*device_pool_); + return *device_context; } absl::StatusOr SyclDevicePool::GetDeviceCount() { ABSL_RETURN_IF_ERROR(SyclDevicePool::InitDevicePool()); // Cast to int since device_ordinal is usually an int. - return static_cast(device_pool_.size()); + return static_cast(device_pool_->size()); } absl::StatusOr SyclDevicePool::GetDeviceOrdinal( const ::sycl::device& device) { ABSL_RETURN_IF_ERROR(SyclDevicePool::InitDevicePool()); - auto it = std::find(device_pool_.begin(), device_pool_.end(), device); - if (it != device_pool_.end()) { - return static_cast(it - device_pool_.begin()); + auto it = std::find(device_pool_->begin(), device_pool_->end(), device); + if (it != device_pool_->end()) { + return static_cast(it - device_pool_->begin()); } return absl::InternalError( "SyclDevicePool::GetDeviceOrdinal failed, got invalid device"); @@ -249,10 +252,10 @@ absl::StatusOr<::sycl::device> SyclDevicePool::GetDevice(int device_ordinal) { ABSL_RETURN_IF_ERROR(SyclDevicePool::InitDevicePool()); ABSL_RETURN_IF_ERROR( IsValidDeviceOrdinal(device_ordinal, "SyclDevicePool::GetDevice")); - return device_pool_[device_ordinal]; + return (*device_pool_)[device_ordinal]; } -StreamPoolMap SyclStreamPool::stream_pool_map_; +absl::NoDestructor SyclStreamPool::stream_pool_map_; absl::Mutex SyclStreamPool::stream_pool_mu_(absl::kConstInit); void SyclAsyncHandler(::sycl::exception_list ex_list) { @@ -269,10 +272,10 @@ void SyclAsyncHandler(::sycl::exception_list ex_list) { absl::StatusOr SyclStreamPool::InitStreamPool(int device_ordinal) { { absl::ReaderMutexLock read_lock(&stream_pool_mu_); - auto it = stream_pool_map_.find(device_ordinal); + auto it = stream_pool_map_->find(device_ordinal); // Returns the existing non-empty stream pool for this device, if available. // The pool may be empty if DestroyStream was called on the last stream. - if (it != stream_pool_map_.end() && !it->second.empty()) { + if (it != stream_pool_map_->end() && !it->second.empty()) { VLOG(2) << "Check 1: Returning existing stream pool for device ordinal " << device_ordinal << " whose size is " << it->second.size(); return &(it->second); @@ -283,14 +286,14 @@ absl::StatusOr SyclStreamPool::InitStreamPool(int device_ordinal) { ::sycl::property::queue::in_order()}; ABSL_ASSIGN_OR_RETURN(::sycl::device sycl_device, SyclDevicePool::GetDevice(device_ordinal)); - ABSL_ASSIGN_OR_RETURN(::sycl::context sycl_context, + ABSL_ASSIGN_OR_RETURN(const ::sycl::context& sycl_context, SyclDevicePool::GetDeviceContext()); VLOG(2) << "Creating new stream pool for device ordinal " << device_ordinal; absl::MutexLock write_lock(&stream_pool_mu_); - auto it = stream_pool_map_.find(device_ordinal); + auto it = stream_pool_map_->find(device_ordinal); // Double-checks that another thread has not already created the pool. - if (it != stream_pool_map_.end() && !it->second.empty()) { + if (it != stream_pool_map_->end() && !it->second.empty()) { VLOG(2) << "Check 2: Returning existing stream pool for device ordinal " << device_ordinal << " whose size is " << it->second.size(); return &(it->second); @@ -301,9 +304,9 @@ absl::StatusOr SyclStreamPool::InitStreamPool(int device_ordinal) { // Use assignment (not insert) to update the stream pool if it was // previously destroyed. - stream_pool_map_[device_ordinal] = std::move(stream_pool); + (*stream_pool_map_)[device_ordinal] = std::move(stream_pool); - return &(stream_pool_map_[device_ordinal]); + return &((*stream_pool_map_)[device_ordinal]); } absl::StatusOr SyclStreamPool::GetDefaultStream(int device_ordinal) { @@ -344,7 +347,7 @@ absl::StatusOr SyclStreamPool::GetOrCreateStream( ::sycl::property::queue::in_order()}; ABSL_ASSIGN_OR_RETURN(::sycl::device sycl_device, SyclDevicePool::GetDevice(device_ordinal)); - ABSL_ASSIGN_OR_RETURN(::sycl::context sycl_context, + ABSL_ASSIGN_OR_RETURN(const ::sycl::context& sycl_context, SyclDevicePool::GetDeviceContext()); stream_pool->push_back(std::make_shared<::sycl::queue>( sycl_context, sycl_device, SyclAsyncHandler, prop_list)); @@ -406,7 +409,7 @@ absl::Status SyclStreamPool::DestroyStream(int device_ordinal, void SyclStreamPool::Reset() { absl::MutexLock write_lock(&stream_pool_mu_); - for (auto& [device_ordinal, stream_pool] : stream_pool_map_) { + for (auto& [device_ordinal, stream_pool] : *stream_pool_map_) { for (auto& stream_handle : stream_pool) { if (stream_handle) { stream_handle->wait(); @@ -415,7 +418,7 @@ void SyclStreamPool::Reset() { } stream_pool.clear(); } - stream_pool_map_.clear(); + stream_pool_map_->clear(); } absl::StatusOr SyclGetTimerProperties(int device_ordinal) { diff --git a/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime.h b/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime.h index 0c3fcc76e3724b..ce0814957f7e04 100644 --- a/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime.h +++ b/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime.h @@ -26,6 +26,7 @@ limitations under the License. #include #include "absl/base/attributes.h" +#include "absl/base/no_destructor.h" #include "absl/container/flat_hash_map.h" #include "absl/strings/ascii.h" #include "xla/stream_executor/sycl/sycl_status.h" @@ -53,7 +54,7 @@ class SyclDevicePool { // This function assumes that the device pool is not modified after // initialization. If this assumption is violated, the context may become // invalid. - static absl::StatusOr<::sycl::context> GetDeviceContext(); + static absl::StatusOr GetDeviceContext(); // Returns the number of devices in the pool. static absl::StatusOr GetDeviceCount(); @@ -65,8 +66,11 @@ class SyclDevicePool { static absl::StatusOr<::sycl::device> GetDevice(int device_ordinal); private: - // The underlying device pool. - static DevicePool device_pool_; + // Leaked: libsycl/Level Zero teardown order relative to this static is + // oneAPI-version-dependent (after it pre-2026.x, before it in 2026.x+), + // making its destructor unsafe to run, so cleanup is deferred to the OS + // instead. Acceptable since this is a one-time allocation per program run. + static absl::NoDestructor device_pool_; // Thread-safe initialization of device_pool_ with all Level-Zero backend GPUs // using absl::call_once. @@ -120,8 +124,10 @@ class SyclStreamPool { static absl::Mutex stream_pool_mu_; // The underlying stream pool for each device. The device ordinal - // is used as the key. - static StreamPoolMap stream_pool_map_ ABSL_GUARDED_BY(stream_pool_mu_); + // is used as the key. Leaked for the same reason as + // SyclDevicePool::device_pool_. + static absl::NoDestructor stream_pool_map_ + ABSL_GUARDED_BY(stream_pool_mu_); // Initializes and returns a pointer to the stream pool for the given device // ordinal. diff --git a/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime_test.cc b/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime_test.cc index 1e733a0d18edf4..190a4141539e8a 100644 --- a/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime_test.cc +++ b/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime_test.cc @@ -126,9 +126,9 @@ TEST_F(SyclGpuRuntimeTest, GetDeviceOrdinal) { TEST_F(SyclGpuRuntimeTest, TestStaticDeviceContext) { // Verify that GetDeviceContext returns the same context instance on multiple // calls. - TF_ASSERT_OK_AND_ASSIGN(::sycl::context saved_sycl_context, + TF_ASSERT_OK_AND_ASSIGN(const ::sycl::context& saved_sycl_context, SyclDevicePool::GetDeviceContext()); - TF_ASSERT_OK_AND_ASSIGN(::sycl::context current_sycl_context, + TF_ASSERT_OK_AND_ASSIGN(const ::sycl::context& current_sycl_context, SyclDevicePool::GetDeviceContext()); EXPECT_EQ(saved_sycl_context, current_sycl_context); } diff --git a/third_party/xla/xla/tools/xla_compile_lib.cc b/third_party/xla/xla/tools/xla_compile_lib.cc index 176f4819ecd522..8f5a59d0e96aa4 100644 --- a/third_party/xla/xla/tools/xla_compile_lib.cc +++ b/third_party/xla/xla/tools/xla_compile_lib.cc @@ -333,6 +333,7 @@ absl::StatusOr LoadAutotuneDataFromModule(HloModuleAndMetadata* mod, if (auto* data = static_cast( mod->backend_specific_data.get()); data != nullptr && data->autotune_results.has_value() && + !data->autotune_results->results().empty() && mod->hlo_module->config().debug_options().xla_gpu_autotune_level() > 0) { ABSL_RETURN_IF_ERROR( diff --git a/third_party/xla/xla/tools/xla_gpu_compile_lib_test.cc b/third_party/xla/xla/tools/xla_gpu_compile_lib_test.cc index 03ce87a90a5eff..c850c02df1b3f9 100644 --- a/third_party/xla/xla/tools/xla_gpu_compile_lib_test.cc +++ b/third_party/xla/xla/tools/xla_gpu_compile_lib_test.cc @@ -263,6 +263,25 @@ TEST_F(XlaCompileLibTest, EXPECT_TRUE(gpu::AutotunerCache::ResultCacheIsEmpty()); } +TEST_F(XlaCompileLibTest, + LoadAutotuneDataGpuDataPresentEmptyAndAutotuningEnabled) { + gpu::AutotunerCache::ClearAutotuneResults(); + + HloModuleAndMetadata mod; + mod.hlo_module = std::move(module_); + auto data = std::make_unique(); + data->autotune_results.emplace(); // empty results + mod.backend_specific_data = std::move(data); + + DebugOptions opts = mod.hlo_module->config().debug_options(); + opts.set_xla_gpu_autotune_level(3); + mod.hlo_module->mutable_config().set_debug_options(opts); + + EXPECT_THAT(internal::LoadAutotuneDataFromModule(&mod, BackendType::kGpu), + absl_testing::IsOkAndHolds(false)); + EXPECT_TRUE(gpu::AutotunerCache::ResultCacheIsEmpty()); +} + TEST_F(XlaCompileLibTest, MainForGpuForceAutoLayout) { static constexpr absl::string_view kHloText = R"( HloModule f