From 5d630bdf9140046c03115f29436bea066f4481ea Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 10 Mar 2026 13:37:28 +0000 Subject: [PATCH 01/16] Bump markdown in /ci/official/requirements_updater/numpy1_requirements Bumps [markdown](https://github.com/Python-Markdown/markdown) from 3.8 to 3.8.1. - [Release notes](https://github.com/Python-Markdown/markdown/releases) - [Changelog](https://github.com/Python-Markdown/markdown/blob/master/docs/changelog.md) - [Commits](https://github.com/Python-Markdown/markdown/compare/3.8...3.8.1) --- updated-dependencies: - dependency-name: markdown dependency-version: 3.8.1 dependency-type: direct:production ... Signed-off-by: dependabot[bot] --- .../numpy1_requirements/requirements_lock_3_10.txt | 6 +++--- .../numpy1_requirements/requirements_lock_3_11.txt | 6 +++--- .../numpy1_requirements/requirements_lock_3_12.txt | 6 +++--- 3 files changed, 9 insertions(+), 9 deletions(-) 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 4767ae65e97159..158308df373f9f 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 @@ -248,9 +248,9 @@ libclang==18.1.1 \ lit==17.0.6 \ --hash=sha256:dfa9af9b55fc4509a56be7bf2346f079d7f4a242d583b9f2e0b078fd0abae31b # via -r ci/official/requirements_updater/requirements.in -markdown==3.8 \ - --hash=sha256:794a929b79c5af141ef5ab0f2f642d0f7b1872981250230e72682346f7cc90dc \ - --hash=sha256:7df81e63f0df5c4b24b7d156eb81e4690595239b7d70937d0409f1b0de319c6f +markdown==3.8.1 \ + --hash=sha256:46cc0c0f1e5211ab2e9d453582f0b28a1bfaf058a9f7d5c50386b99b588d8811 \ + --hash=sha256:a2e2f01cead4828ee74ecca9623045f62216aef2212a7685d6eb9163f590b8c1 # via tb-nightly markdown-it-py==3.0.0 \ --hash=sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1 \ 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 a7a2e49d19bfee..00718cd4ae243d 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 @@ -248,9 +248,9 @@ libclang==18.1.1 \ lit==17.0.6 \ --hash=sha256:dfa9af9b55fc4509a56be7bf2346f079d7f4a242d583b9f2e0b078fd0abae31b # via -r ci/official/requirements_updater/requirements.in -markdown==3.8 \ - --hash=sha256:794a929b79c5af141ef5ab0f2f642d0f7b1872981250230e72682346f7cc90dc \ - --hash=sha256:7df81e63f0df5c4b24b7d156eb81e4690595239b7d70937d0409f1b0de319c6f +markdown==3.8.1 \ + --hash=sha256:46cc0c0f1e5211ab2e9d453582f0b28a1bfaf058a9f7d5c50386b99b588d8811 \ + --hash=sha256:a2e2f01cead4828ee74ecca9623045f62216aef2212a7685d6eb9163f590b8c1 # via tb-nightly markdown-it-py==3.0.0 \ --hash=sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1 \ 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 ba5a66481aece5..60bb3481ad11a7 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 @@ -248,9 +248,9 @@ libclang==18.1.1 \ lit==17.0.6 \ --hash=sha256:dfa9af9b55fc4509a56be7bf2346f079d7f4a242d583b9f2e0b078fd0abae31b # via -r ci/official/requirements_updater/requirements.in -markdown==3.8 \ - --hash=sha256:794a929b79c5af141ef5ab0f2f642d0f7b1872981250230e72682346f7cc90dc \ - --hash=sha256:7df81e63f0df5c4b24b7d156eb81e4690595239b7d70937d0409f1b0de319c6f +markdown==3.8.1 \ + --hash=sha256:46cc0c0f1e5211ab2e9d453582f0b28a1bfaf058a9f7d5c50386b99b588d8811 \ + --hash=sha256:a2e2f01cead4828ee74ecca9623045f62216aef2212a7685d6eb9163f590b8c1 # via tb-nightly markdown-it-py==3.0.0 \ --hash=sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1 \ From e923cf41cb9bad2f78aacb3539cf8224dd572075 Mon Sep 17 00:00:00 2001 From: Huy Phung Date: Sat, 29 Aug 2026 04:05:19 -0700 Subject: [PATCH 02/16] fix segfault by adding initialize for deletion --- tensorflow/core/data/service/server_lib.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/data/service/server_lib.h b/tensorflow/core/data/service/server_lib.h index 56a8f8d94fc558..8647da848591ea 100644 --- a/tensorflow/core/data/service/server_lib.h +++ b/tensorflow/core/data/service/server_lib.h @@ -124,7 +124,7 @@ class DispatchGrpcDataServer : public GrpcDataServerBase { private: const experimental::DispatcherConfig config_; // Owned. We use a raw pointer because GrpcDispatcherImpl is forward-declared. - GrpcDispatcherImpl* service_; + GrpcDispatcherImpl* service_ = nullptr; }; // A wrapper for `SnapshotTaskProgress` for use with pybind. @@ -170,7 +170,7 @@ class WorkerGrpcDataServer : public GrpcDataServerBase { const experimental::WorkerConfig config_; // Owned. We use a raw pointer because GrpcWorkerImpl is forward-declared. - GrpcWorkerImpl* service_; + GrpcWorkerImpl* service_ = nullptr; std::shared_ptr transfer_server_; }; From 52878813fecfe0d718cf3a85c2e2e1723a754139 Mon Sep 17 00:00:00 2001 From: Huy Phung Date: Sat, 29 Aug 2026 22:49:55 -0700 Subject: [PATCH 03/16] fix the join() function to not waiting when start = false|none --- tensorflow/core/data/service/server_lib.cc | 12 ++++++++++-- tensorflow/core/data/service/server_lib.h | 2 +- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/tensorflow/core/data/service/server_lib.cc b/tensorflow/core/data/service/server_lib.cc index e0b5f3cf9ac90b..4f23427e7c2891 100644 --- a/tensorflow/core/data/service/server_lib.cc +++ b/tensorflow/core/data/service/server_lib.cc @@ -21,9 +21,9 @@ limitations under the License. #include #include +#include "absl/strings/str_cat.h" #include "grpcpp/server.h" #include "grpcpp/server_builder.h" -#include "absl/strings/str_cat.h" #include "tensorflow/core/data/service/common.pb.h" #include "tensorflow/core/data/service/credentials_factory.h" #include "tensorflow/core/data/service/export.pb.h" @@ -103,7 +103,15 @@ void GrpcDataServerBase::Stop() { stopped_ = true; } -void GrpcDataServerBase::Join() { server_->Wait(); } +absl::Status GrpcDataServerBase::Join() { + if (!server_) { + return absl::FailedPreconditionError( + "Server cannot be joined before it has been started. " + "Call `start()` before `join()`."); + } + server_->Wait(); + return absl::OkStatus(); +} int GrpcDataServerBase::BoundPort() { return bound_port(); } diff --git a/tensorflow/core/data/service/server_lib.h b/tensorflow/core/data/service/server_lib.h index 8647da848591ea..e0c0f0ccbdeda8 100644 --- a/tensorflow/core/data/service/server_lib.h +++ b/tensorflow/core/data/service/server_lib.h @@ -56,7 +56,7 @@ class GrpcDataServerBase { void Stop(); // Blocks until the server stops. - void Join(); + absl::Status Join(); // Returns the port bound by the server. Only valid after calling Start(). int BoundPort(); From b08bc3dd8e4facef024c5503b5d4e7a8c814c308 Mon Sep 17 00:00:00 2001 From: Huy Phung Date: Sat, 29 Aug 2026 22:50:48 -0700 Subject: [PATCH 04/16] add testcase when worker or dispatcher start failed combine with join() function --- .../experimental/service/server_lib_test.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tensorflow/python/data/experimental/service/server_lib_test.py b/tensorflow/python/data/experimental/service/server_lib_test.py index 6d44747def1668..66d0af0b4c0655 100644 --- a/tensorflow/python/data/experimental/service/server_lib_test.py +++ b/tensorflow/python/data/experimental/service/server_lib_test.py @@ -148,6 +148,30 @@ def testJoinWorker(self): worker.stop() worker.join() + def testJoinDispatcherWithoutStart(self): + dispatcher = server_lib.DispatchServer(start=False) + with self.assertRaisesRegex( + RuntimeError, "Server cannot be joined before it has been started"): + dispatcher.join() + + def testJoinWorkerWithoutStart(self): + dispatcher = server_lib.DispatchServer() + worker = server_lib.WorkerServer( + server_lib.WorkerConfig(dispatcher._address), start=False) + with self.assertRaisesRegex( + RuntimeError, "Server cannot be joined before it has been started"): + worker.join() + + def testDispatcherDestroyedWithoutStart(self): + dispatcher = server_lib.DispatchServer(start=False) + del dispatcher + + def testWorkerDestroyedWithoutStart(self): + dispatcher = server_lib.DispatchServer() + worker = server_lib.WorkerServer( + server_lib.WorkerConfig(dispatcher._address), start=False) + del worker + def testDispatcherNumWorkers(self): dispatcher = server_lib.DispatchServer() self.assertEqual(0, dispatcher._num_workers()) From f1e8404afb9d52f92833b6f748a2d8b3a6946a65 Mon Sep 17 00:00:00 2001 From: Huy Phung Date: Mon, 31 Aug 2026 00:56:50 -0700 Subject: [PATCH 05/16] ci: trigger rebuild From 0fae1d12bf9e64dadebd8417f17357b86a3f094c Mon Sep 17 00:00:00 2001 From: Huy Phung Date: Mon, 31 Aug 2026 22:38:26 -0700 Subject: [PATCH 06/16] revert join function back to void --- tensorflow/core/data/service/server_lib.cc | 7 ++----- tensorflow/core/data/service/server_lib.h | 2 +- .../data/experimental/service/server_lib_test.py | 11 ++++------- 3 files changed, 7 insertions(+), 13 deletions(-) diff --git a/tensorflow/core/data/service/server_lib.cc b/tensorflow/core/data/service/server_lib.cc index 4f23427e7c2891..2600e5a15e9e52 100644 --- a/tensorflow/core/data/service/server_lib.cc +++ b/tensorflow/core/data/service/server_lib.cc @@ -103,14 +103,11 @@ void GrpcDataServerBase::Stop() { stopped_ = true; } -absl::Status GrpcDataServerBase::Join() { +void GrpcDataServerBase::Join() { if (!server_) { - return absl::FailedPreconditionError( - "Server cannot be joined before it has been started. " - "Call `start()` before `join()`."); + return; } server_->Wait(); - return absl::OkStatus(); } int GrpcDataServerBase::BoundPort() { return bound_port(); } diff --git a/tensorflow/core/data/service/server_lib.h b/tensorflow/core/data/service/server_lib.h index e0c0f0ccbdeda8..8647da848591ea 100644 --- a/tensorflow/core/data/service/server_lib.h +++ b/tensorflow/core/data/service/server_lib.h @@ -56,7 +56,7 @@ class GrpcDataServerBase { void Stop(); // Blocks until the server stops. - absl::Status Join(); + void Join(); // Returns the port bound by the server. Only valid after calling Start(). int BoundPort(); diff --git a/tensorflow/python/data/experimental/service/server_lib_test.py b/tensorflow/python/data/experimental/service/server_lib_test.py index 66d0af0b4c0655..f195bf289e5187 100644 --- a/tensorflow/python/data/experimental/service/server_lib_test.py +++ b/tensorflow/python/data/experimental/service/server_lib_test.py @@ -150,17 +150,14 @@ def testJoinWorker(self): def testJoinDispatcherWithoutStart(self): dispatcher = server_lib.DispatchServer(start=False) - with self.assertRaisesRegex( - RuntimeError, "Server cannot be joined before it has been started"): - dispatcher.join() + dispatcher.join() def testJoinWorkerWithoutStart(self): dispatcher = server_lib.DispatchServer() worker = server_lib.WorkerServer( - server_lib.WorkerConfig(dispatcher._address), start=False) - with self.assertRaisesRegex( - RuntimeError, "Server cannot be joined before it has been started"): - worker.join() + server_lib.WorkerConfig(dispatcher._address), start=False + ) + worker.join() def testDispatcherDestroyedWithoutStart(self): dispatcher = server_lib.DispatchServer(start=False) From 9ddd5014516e1ad2ad1098de99e75d315e8d534e Mon Sep 17 00:00:00 2001 From: Levon Ter-Grigoryan Date: Tue, 1 Sep 2026 05:38:30 -0700 Subject: [PATCH 07/16] [XLA:GPU] Rename collective_kernel_api to multi_gpu_barrier. This API is not used by custom kernels anymore. Right now it contains only cross-device barrier related logic. PiperOrigin-RevId: 974442543 --- third_party/xla/xla/backends/gpu/collectives/BUILD | 2 +- .../backends/gpu/collectives/nccl_communicator.cc | 2 +- third_party/xla/xla/backends/gpu/runtime/BUILD | 11 +++++------ .../xla/backends/gpu/runtime/collective_cliques.cc | 2 +- ...collective_kernel_api.cc => multi_gpu_barrier.cc} | 2 +- .../{collective_kernel_api.h => multi_gpu_barrier.h} | 12 +++--------- .../backends/gpu/runtime/ragged_all_to_all_thunk.cc | 2 +- 7 files changed, 13 insertions(+), 20 deletions(-) rename third_party/xla/xla/backends/gpu/runtime/{collective_kernel_api.cc => multi_gpu_barrier.cc} (98%) rename third_party/xla/xla/backends/gpu/runtime/{collective_kernel_api.h => multi_gpu_barrier.h} (81%) diff --git a/third_party/xla/xla/backends/gpu/collectives/BUILD b/third_party/xla/xla/backends/gpu/collectives/BUILD index 09c9c6c0e3f109..6e6aadc1c6ff5e 100644 --- a/third_party/xla/xla/backends/gpu/collectives/BUILD +++ b/third_party/xla/xla/backends/gpu/collectives/BUILD @@ -1103,7 +1103,7 @@ cc_library( "//xla:shape_util", "//xla:util", "//xla:xla_data_proto_cc", - "//xla/backends/gpu/runtime:collective_kernel_api", + "//xla/backends/gpu/runtime:multi_gpu_barrier", "//xla/core/collectives:communicator", "//xla/core/collectives:rank_id", "//xla/core/collectives:reduction_kind", diff --git a/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.cc b/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.cc index 43d9f3b7073205..cad31b7fe963e4 100644 --- a/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.cc +++ b/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.cc @@ -49,7 +49,7 @@ limitations under the License. #include "xla/backends/gpu/collectives/nccl_symmetric_memory.h" #include "xla/backends/gpu/collectives/nccl_types.h" #include "xla/backends/gpu/collectives/single_threaded_executor.h" -#include "xla/backends/gpu/runtime/collective_kernel_api.h" +#include "xla/backends/gpu/runtime/multi_gpu_barrier.h" #include "xla/core/collectives/communicator.h" #include "xla/core/collectives/rank_id.h" #include "xla/core/collectives/reduction_kind.h" diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD index 75bdcc99a3ef38..53b24c32bff7be 100644 --- a/third_party/xla/xla/backends/gpu/runtime/BUILD +++ b/third_party/xla/xla/backends/gpu/runtime/BUILD @@ -2360,12 +2360,12 @@ cc_library( deps = [ ":collective_clique_requests", ":collective_cliques", - ":collective_kernel_api", ":collective_memory", ":collective_memory_requests", ":collective_thunk", ":collective_thunk_proto_cc", ":command_state", + ":multi_gpu_barrier", ":ragged_all_to_all", ":thunk", ":thunk_proto_cc", @@ -2946,8 +2946,8 @@ cc_library( compatible_with = get_compatible_with_portable(), deps = [ ":collective_clique_requests", - ":collective_kernel_api", ":collective_params", + ":multi_gpu_barrier", "//xla:status_macros", "//xla:util", "//xla/backends/gpu/collectives:gpu_clique", @@ -3431,9 +3431,9 @@ cc_library( ) cc_library( - name = "collective_kernel_api", - srcs = ["collective_kernel_api.cc"], - hdrs = ["collective_kernel_api.h"], + name = "multi_gpu_barrier", + srcs = ["multi_gpu_barrier.cc"], + hdrs = ["multi_gpu_barrier.h"], compatible_with = get_compatible_with_portable(), deps = [ "//xla:status_macros", @@ -3449,7 +3449,6 @@ cc_library( "@com_google_absl//absl/base:no_destructor", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/status", - "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/synchronization", ], diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_cliques.cc b/third_party/xla/xla/backends/gpu/runtime/collective_cliques.cc index 3855c93ee98768..b0fb8788dc98d8 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_cliques.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_cliques.cc @@ -36,8 +36,8 @@ limitations under the License. #include "xla/backends/gpu/collectives/gpu_cliques.h" #include "xla/backends/gpu/collectives/gpu_communicator.h" #include "xla/backends/gpu/runtime/collective_clique_requests.h" -#include "xla/backends/gpu/runtime/collective_kernel_api.h" #include "xla/backends/gpu/runtime/collective_params.h" +#include "xla/backends/gpu/runtime/multi_gpu_barrier.h" #include "xla/core/collectives/clique_id.h" #include "xla/core/collectives/clique_key.h" #include "xla/core/collectives/communicator.h" diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc b/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc similarity index 98% rename from third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc rename to third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc index 0eb6811f600d7b..4f2ce8c8f39aa1 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc +++ b/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/backends/gpu/runtime/collective_kernel_api.h" +#include "xla/backends/gpu/runtime/multi_gpu_barrier.h" #include #include diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h b/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.h similarity index 81% rename from third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h rename to third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.h index cd601d10610dd9..c38710e17e15ff 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h +++ b/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_KERNEL_API_H_ -#define XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_KERNEL_API_H_ +#ifndef XLA_BACKENDS_GPU_RUNTIME_MULTI_GPU_BARRIER_H_ +#define XLA_BACKENDS_GPU_RUNTIME_MULTI_GPU_BARRIER_H_ #include #include @@ -28,12 +28,6 @@ limitations under the License. namespace xla::gpu { -// This file contains collective kernels related method shared between XLA and -// Mosaic. Be careful when changing the API. -// -// Collective kernels are kernels which are launched simultaneously on multiple -// devices, and have an access to each other device's memory. - // Launches a cross-GPU barrier synchronization. absl::Status LaunchMultiGpuBarrier( stream_executor::Stream* stream, int64_t num_devices, RankId rank, @@ -57,4 +51,4 @@ size_t GetMultiGpuBarrierSignalValueSize(); } // namespace xla::gpu -#endif // XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_KERNEL_API_H_ +#endif // XLA_BACKENDS_GPU_RUNTIME_MULTI_GPU_BARRIER_H_ diff --git a/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc index 547dea71fffb9b..4c4b60f7cf0277 100644 --- a/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc @@ -48,12 +48,12 @@ limitations under the License. #include "xla/backends/gpu/collectives/gxl_communicator.h" #include "xla/backends/gpu/runtime/collective_clique_requests.h" #include "xla/backends/gpu/runtime/collective_cliques.h" -#include "xla/backends/gpu/runtime/collective_kernel_api.h" #include "xla/backends/gpu/runtime/collective_memory.h" #include "xla/backends/gpu/runtime/collective_memory_requests.h" #include "xla/backends/gpu/runtime/collective_thunk.h" #include "xla/backends/gpu/runtime/collective_thunk.pb.h" #include "xla/backends/gpu/runtime/command_state.h" +#include "xla/backends/gpu/runtime/multi_gpu_barrier.h" #include "xla/backends/gpu/runtime/ragged_all_to_all.h" #include "xla/backends/gpu/runtime/thunk.h" #include "xla/backends/gpu/runtime/thunk.pb.h" From e26cbf5b2bc6a4d56e2bed8bbe24d335b42306ce Mon Sep 17 00:00:00 2001 From: Tori Baker Date: Tue, 1 Sep 2026 05:47:59 -0700 Subject: [PATCH 08/16] Disable experimental GEMM fusion v2 in HoistFusedBitcastsReshapeTest. Overrides GetDebugOptionsForTest to set xla_gpu_experimental_gemm_fusion_v2 to false for HoistFusedBitcastsReshapeTest. We don't run HoistFusedBitcasts when gemm_fusion_v2=true, so this test doesn't make sense for it. Set it to false explicitely. PiperOrigin-RevId: 974445800 --- .../backends/gpu/transforms/hoist_fused_bitcasts_test.cc | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/third_party/xla/xla/backends/gpu/transforms/hoist_fused_bitcasts_test.cc b/third_party/xla/xla/backends/gpu/transforms/hoist_fused_bitcasts_test.cc index 9bec365d533c46..ddbd4ca0ec34e4 100644 --- a/third_party/xla/xla/backends/gpu/transforms/hoist_fused_bitcasts_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/hoist_fused_bitcasts_test.cc @@ -76,6 +76,14 @@ class HoistFusedBitcastsReshapeTest HoistFusedBitcastsReshapeTest() { RegisterSymbolicExprStorage(&mlir_context_); } + + DebugOptions GetDebugOptionsForTest() const override { + DebugOptions debug_options = + HloHardwareIndependentTestBase::GetDebugOptionsForTest(); + debug_options.set_xla_gpu_experimental_gemm_fusion_v2(false); + return debug_options; + } + const se::DeviceDescription device_description_{ TestGpuDeviceInfo::RTXA6000DeviceInfo( se::GpuComputeCapability{se::CudaComputeCapability::Ampere()})}; From e3b7ac8bf538139dabb7636ed2c0c0dc09248af0 Mon Sep 17 00:00:00 2001 From: Dirk Hornung Date: Tue, 1 Sep 2026 06:59:00 -0700 Subject: [PATCH 09/16] [XLA:GPU] Support duplicating cheap operands in cuDNN conv fusion. This reduces the number of kernels from 4 to 2 in the following example: ``` // Input HLO: %p = f16[32] parameter(0) %b = f16[1,9,9,32] broadcast(%p), dimensions={3} %add1 = add(%conv1, %b) %add2 = add(%conv2, %b) // Before (4 separate kernels emitted): %conv1_out = convolution(...) %add1_out = fusion(%conv1_out, %b), kind=kLoop %conv2_out = convolution(...) %add2_out = fusion(%conv2_out, %b), kind=kLoop // After (2 fused cuDNN graph kernels, 0 intermediate HBM traffic): %fusion1 = fusion(%input1, %filter1, %p), kind=kCustom, calls={ %p0 = f16[32] parameter(2) %b0 = f16[1,9,9,32] broadcast(%p0), dimensions={3} ROOT %out = add(%conv1, %b0) } %fusion2 = fusion(%input2, %filter2, %p), kind=kCustom, calls={ %p0 = f16[32] parameter(2) %b0 = f16[1,9,9,32] broadcast(%p0), dimensions={3} ROOT %out = add(%conv2, %b0) } ``` PiperOrigin-RevId: 974473390 --- .../gpu/transforms/conv_fusion_rewriter.cc | 9 +++ .../transforms/conv_fusion_rewriter_test.cc | 32 +++++++++ .../gpu/transforms/cudnn_fusion_compiler.cc | 69 +++++++++++++++---- .../gpu/transforms/cudnn_fusion_utils.cc | 21 +++++- 4 files changed, 116 insertions(+), 15 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter.cc index 7cef77dadcac8a..7022d28feceac9 100644 --- a/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter.cc +++ b/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter.cc @@ -82,6 +82,15 @@ bool IsConvFusionOutputsValid(const std::vector& outputs) { (outputs[1]->opcode() == HloOpcode::kReduce)) { return false; } + // Disallow upcast converts at fusion outputs. S32->F32 is allowed as cuDNN + // INT8 convs require it for epilogue fusions. + for (const HloInstruction* output : outputs) { + if (output->opcode() == HloOpcode::kConvert && + output->shape().element_type() == F32 && + output->operand(0)->shape().element_type() != S32) { + return false; + } + } return true; } diff --git a/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc index 7cf3fa00701406..ac6e739d45f2d4 100644 --- a/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc @@ -575,6 +575,38 @@ TEST_F(ConvFusionRewriterUnitTest, Test1DBiasBroadcastFusedF16) { .WithShape(F16, {1, 9, 9, 32})); } +TEST_F(ConvFusionRewriterUnitTest, Test1DBiasBroadcastSharedWithMultipleConvs) { + RunAndMatch( + R"( + HloModule Test + + ENTRY Test { + input1 = f16[1,9,9,17] parameter(0) + filter1 = f16[32,3,3,17] parameter(1) + filter2 = f16[32,3,3,32] parameter(2) + bias = f16[32] parameter(3) + bias_broadcast = f16[1,9,9,32] broadcast(bias), dimensions={3} + zero = f16[] constant(0) + zeros = f16[1,9,9,32] broadcast(zero), dimensions={} + + conv1 = f16[1,9,9,32] convolution(input1, filter1), + window={size=3x3 pad=1_1x1_1}, + dim_labels=b01f_o01i->b01f + sum1 = add(conv1, bias_broadcast) + relu1 = maximum(sum1, zeros) + + conv2 = f16[1,9,9,32] convolution(relu1, filter2), + window={size=3x3 pad=1_1x1_1}, + dim_labels=b01f_o01i->b01f + sum2 = add(conv2, bias_broadcast) + ROOT relu2 = maximum(sum2, zeros) + })", + m::Fusion(m::Fusion(m::Parameter(0), m::Parameter(1), m::Parameter(3)), + m::Parameter(2), m::Parameter(3)) + .WithFusionKind(HloInstruction::FusionKind::kCustom) + .WithShape(F16, {1, 9, 9, 32})); +} + TEST_F(ConvFusionRewriterUnitTest, FuseAlpha) { MAYBE_SKIP_TEST("I8"); RunAndMatch(R"( diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler.cc index d20b178c722844..3475a14ab92ef8 100644 --- a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler.cc +++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler.cc @@ -529,6 +529,31 @@ class ConvDimensionAdapter { return -1; } + const HloInstruction* get_broadcast_user(const HloInstruction* hlo) { + auto all_users_are_broadcast = [](const HloInstruction* instr) { + return !instr->users().empty() && + absl::c_all_of(instr->users(), [](const HloInstruction* u) { + return u->opcode() == HloOpcode::kBroadcast; + }); + }; + + // Pattern 1: hlo -> broadcast + if (all_users_are_broadcast(hlo)) { + return hlo->users()[0]; + } + + // Pattern 2: hlo -> convert -> broadcast + if (hlo->user_count() == 1 && + hlo->users()[0]->opcode() == HloOpcode::kConvert) { + const HloInstruction* convert = hlo->users()[0]; + if (all_users_are_broadcast(convert)) { + return convert->users()[0]; + } + } + + return nullptr; + }; + std::optional DimensionsAndStrides(const HloInstruction& hlo) { int64_t spatial_dims = std::max(2, dums_.input_spatial_dimensions_size()); @@ -549,9 +574,8 @@ class ConvDimensionAdapter { // If the parameter is consumed by a broadcast, map its dimensions to the // corresponding cuDNN canonical axes (N, C, spatial...). - if (hlo.user_count() == 1 && - hlo.users()[0]->opcode() == HloOpcode::kBroadcast) { - const auto& bcast_dims = hlo.users()[0]->dimensions(); + if (const HloInstruction* broadcast = get_broadcast_user(&hlo)) { + const auto& bcast_dims = broadcast->dimensions(); for (int i = 0; i < bcast_dims.size(); ++i) { int64_t cudnn_dim = HloDimToCudnnDim(bcast_dims[i]); if (cudnn_dim >= 0 && cudnn_dim < cudnn_rank) { @@ -559,15 +583,18 @@ class ConvDimensionAdapter { result.strides[cudnn_dim] = 1; } } - } else if (hlo.shape().dimensions().size() == 1) { - // Fallback for un-broadcasted 1D parameters: assume channel bias [1, C, - // 1, 1]. + return result; + } + + // Fallback for un-broadcasted 1D parameters: assume channel bias [1, C, + // 1, 1]. + if (hlo.shape().dimensions().size() == 1) { result.sizes[1] = hlo.shape().dimensions(0); result.strides[1] = 1; - } else { - return std::nullopt; + return result; } - return result; + + return std::nullopt; } // Placeholder FP32 data type here, it is not used. auto desc = se::dnn::TensorDescriptor::For( @@ -831,10 +858,14 @@ absl::StatusOr HloFusionToCuDnnGraph( // and int32 = conv(int8, int8) hlo_to_cudnn[hlo] = operand(0); continue; - } else if (HloPredicateIsOp(hlo)) { + } + + if (HloPredicateIsOp(hlo)) { CHECK(hlo_to_cudnn.contains(hlo)); continue; - } else if (HloPredicateIsOp(hlo)) { + } + + if (HloPredicateIsOp(hlo)) { if (hlo->user_count() != 1 || !IsWorkspaceAllocationRoot(*hlo->users()[0])) { return absl::UnimplementedError( @@ -843,7 +874,9 @@ absl::StatusOr HloFusionToCuDnnGraph( hlo->ToString())); } continue; - } else if (HloPredicateIsOp(hlo)) { + } + + if (HloPredicateIsOp(hlo)) { if (!IsWorkspaceAllocationRoot(*hlo) && !IsAmaxRoot(*hlo)) { return absl::UnimplementedError( absl::StrCat("Tuples are only expected at outputs for workspace " @@ -851,9 +884,17 @@ absl::StatusOr HloFusionToCuDnnGraph( hlo->ToString())); } continue; - } else if (HloPredicateIsOp(hlo)) { + } + + if (HloPredicateIsOp(hlo)) { + const Shape& root_shape = computation.root_instruction()->shape(); + const Shape& output_shape = + root_shape.IsTuple() ? root_shape.tuple_shapes(0) : root_shape; + int64_t rank = std::max(3, output_shape.dimensions().size()); + ABSL_ASSIGN_OR_RETURN(hlo_to_cudnn[hlo], + HandleConstantHloToCudnnGraph(*hlo, graph, rank)); ABSL_ASSIGN_OR_RETURN(hlo_to_cudnn[hlo], - HandleConstantHloToCudnnGraph(*hlo, graph)); + HandleConstantHloToCudnnGraph(*hlo, graph, rank)); } else if (HloPredicateIsOp(hlo)) { diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_utils.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_utils.cc index a60ab8e6ad60fb..dd4b4d18606f3b 100644 --- a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_utils.cc +++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_utils.cc @@ -151,6 +151,25 @@ bool IsEpilogueOpSupportedByCuDNN(const HloInstruction& hlo, } } +namespace { + +bool IsCheapToDuplicate(const HloInstruction* hlo) { + if (hlo->opcode() == HloOpcode::kConstant) { + return ShapeUtil::IsScalar(hlo->shape()); + } + if (hlo->opcode() == HloOpcode::kBroadcast) { + return ShapeUtil::IsScalar(hlo->operand(0)->shape()) || + hlo->operand(0)->shape().dimensions().size() == 1; + } + if (hlo->opcode() == HloOpcode::kConvert) { + return ShapeUtil::IsScalar(hlo->shape()) || + hlo->shape().dimensions().size() == 1; + } + return false; +} + +} // namespace + HloInstruction* FuseTowardOperand( HloInstruction* hlo, HloComputation::Builder& builder, std::vector& fusion_params, @@ -162,7 +181,7 @@ HloInstruction* FuseTowardOperand( HloInstruction* fused_hlo; if (IsEpilogueOpSupportedByCuDNN(*hlo, /*can_fuse_reduce=*/false, is_nchw, device_info) && - hlo->user_count() == 1) { + (hlo->user_count() == 1 || IsCheapToDuplicate(hlo))) { HloInstruction::InstructionVector new_operands; for (int i = 0; i < hlo->operand_count(); ++i) { HloInstruction* operand = hlo->mutable_operand(i); From 91fb2c37526ac8ad9dc14b6e2ffe7c393c2e6f62 Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Tue, 1 Sep 2026 07:35:26 -0700 Subject: [PATCH 10/16] Skip fusion computations in HloComputationDeduplicator. Unless we run with `mark_fusion_duplications = true`, we should just ignore fusion computations. We are not allowed to deduplicate them. PiperOrigin-RevId: 974488637 --- .../simplifiers/hlo_computation_deduplicator.cc | 7 +++++-- .../hlo_computation_deduplicator_test.cc | 14 ++++++++++---- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_computation_deduplicator.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_computation_deduplicator.cc index de2bc2387b4c91..04a6526c803318 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_computation_deduplicator.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_computation_deduplicator.cc @@ -91,9 +91,12 @@ absl::StatusOr HloComputationDeduplicator::RunImpl( module->MakeComputationPostOrder(execution_threads)) { // Ignore entry computation since it is called from outside and computations // with large number of instructions or large-size constants due to increase - // in time taken to stringify. + // in time taken to stringify. Also ignore fusion computations, which need + // to have a 1:1 relationship between caller and computation, except if we + // run in the mode where we want to annotate duplicate fusion computations. if (comp->IsEntryComputation() || comp->instruction_count() > 128 || - ContainsLargeConstants(comp)) { + ContainsLargeConstants(comp) || + (!mark_fusion_duplications_ && comp->IsFusionComputation())) { continue; } // Don't deduplicate collectives and non-collectives. diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_computation_deduplicator_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_computation_deduplicator_test.cc index 9e0b3c860400f9..4b668a3593e09f 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_computation_deduplicator_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_computation_deduplicator_test.cc @@ -21,6 +21,7 @@ limitations under the License. #include #include +#include #include #include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_computation.h" @@ -629,7 +630,7 @@ TEST_F(HloComputationDeduplicatorTest, LargeSubComputationTest) { } module->AddEntryComputation(main.Build()); HloComputationDeduplicator dedup; - TF_ASSERT_OK_AND_ASSIGN(bool changed, dedup.Run(module.get())); + ASSERT_OK_AND_ASSIGN(bool changed, dedup.Run(module.get())); EXPECT_FALSE(changed); std::vector computations = module->MakeComputationSorted(); EXPECT_EQ(computations.size(), (total_regions + 1)); @@ -665,8 +666,8 @@ TEST_F(HloComputationDeduplicatorTest, DontDeduplicateReduceAllReduce) { } TEST_F(HloComputationDeduplicatorTest, DeduplicateChain) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(R"( + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( HloModule module fusion0 { @@ -696,8 +697,13 @@ ENTRY entry { ROOT add = f32[] add(fusion.0, p1) } )")); + HloComputationDeduplicator dedup_do_not_mark( + /*mark_fusion_duplications=*/false); + ASSERT_OK_AND_ASSIGN(bool changed, dedup_do_not_mark.Run(module.get())); + EXPECT_FALSE(changed); + HloComputationDeduplicator dedup(/*mark_fusion_duplications=*/true); - TF_ASSERT_OK_AND_ASSIGN(bool changed, dedup.Run(module.get())); + ASSERT_OK_AND_ASSIGN(changed, dedup.Run(module.get())); EXPECT_TRUE(changed); EXPECT_EQ(module->computation_count(), 4); HloInstruction* fusion0 = From b07caa369586d8d016104ada86ae0ac450b2468c Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Tue, 1 Sep 2026 07:40:29 -0700 Subject: [PATCH 11/16] Automated Code Change PiperOrigin-RevId: 974490864 --- tensorflow/core/kernels/linalg/BUILD | 10 ++++++++-- tensorflow/core/kernels/linalg/lu_op.cc | 2 ++ tensorflow/core/kernels/linalg/matrix_diag_op.cc | 11 +++++++---- tensorflow/core/kernels/linalg/self_adjoint_eig_op.cc | 1 + 4 files changed, 18 insertions(+), 6 deletions(-) diff --git a/tensorflow/core/kernels/linalg/BUILD b/tensorflow/core/kernels/linalg/BUILD index 31a192ee9a5c12..4e16037f225206 100644 --- a/tensorflow/core/kernels/linalg/BUILD +++ b/tensorflow/core/kernels/linalg/BUILD @@ -134,7 +134,10 @@ tf_kernel_library( tf_kernel_library( name = "matrix_diag_op", prefix = "matrix_diag_op", - deps = LINALG_DEPS, + deps = LINALG_DEPS + [ + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", + ], ) tf_kernel_library( @@ -183,7 +186,10 @@ tf_kernel_library( tf_kernel_library( name = "self_adjoint_eig_op", prefix = "self_adjoint_eig_op", - deps = LINALG_DEPS + ["//tensorflow/core:lib_internal"], + deps = LINALG_DEPS + [ + "//tensorflow/core:lib_internal", + "@com_google_absl//absl/status", + ], ) tf_kernel_library( diff --git a/tensorflow/core/kernels/linalg/lu_op.cc b/tensorflow/core/kernels/linalg/lu_op.cc index 998769e31fb0e0..de0f0d2ad1a322 100644 --- a/tensorflow/core/kernels/linalg/lu_op.cc +++ b/tensorflow/core/kernels/linalg/lu_op.cc @@ -17,6 +17,8 @@ limitations under the License. #include #include "absl/container/inlined_vector.h" +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "Eigen/Core" // from @eigen_archive #include "Eigen/LU" // from @eigen_archive #include "tensorflow/core/framework/kernel_def_builder.h" diff --git a/tensorflow/core/kernels/linalg/matrix_diag_op.cc b/tensorflow/core/kernels/linalg/matrix_diag_op.cc index 90fa24b8122ef7..f017da44bd9208 100644 --- a/tensorflow/core/kernels/linalg/matrix_diag_op.cc +++ b/tensorflow/core/kernels/linalg/matrix_diag_op.cc @@ -14,17 +14,19 @@ limitations under the License. ==============================================================================*/ // See docs in ../ops/array_ops.cc. +#include +#include +#include + +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" #define EIGEN_USE_THREADS #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #define EIGEN_USE_GPU #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM -#include "tensorflow/core/kernels/linalg/matrix_diag_op.h" - #include -#include -#include #include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive #include "tensorflow/core/framework/op_kernel.h" @@ -33,6 +35,7 @@ limitations under the License. #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" +#include "tensorflow/core/kernels/linalg/matrix_diag_op.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/threadpool.h" #include "tensorflow/core/platform/logging.h" diff --git a/tensorflow/core/kernels/linalg/self_adjoint_eig_op.cc b/tensorflow/core/kernels/linalg/self_adjoint_eig_op.cc index 816fb97174831f..478d3174e2fca2 100644 --- a/tensorflow/core/kernels/linalg/self_adjoint_eig_op.cc +++ b/tensorflow/core/kernels/linalg/self_adjoint_eig_op.cc @@ -17,6 +17,7 @@ limitations under the License. #include +#include "absl/status/status.h" #include "Eigen/Core" // from @eigen_archive #include "Eigen/Eigenvalues" // from @eigen_archive #include "tensorflow/core/framework/kernel_def_builder.h" From 0de11b4d4c13cc0e0ba56afdce532f9b9a20b9ab Mon Sep 17 00:00:00 2001 From: Tori Baker Date: Tue, 1 Sep 2026 10:06:02 -0700 Subject: [PATCH 12/16] Update TritonGemmFusionTest to be parametrized and test GemmFusionV1 & V2 in preparation of V2 launch. Remove test MultipleBatchRequireSeparateTranspose - this is a fusion heuristic that does not belong here and is different than the new transpose fusion heuristic that we chose for V2. Rewrite LinkingLibdeviceTwiceWorks - this test is meant to confirm that a module with multiple triton gemms is able to launch, but the matcher is too dependent on fusion decisions that are not relevant for the test. As mentioned in the TODO, it would be better if this was rewritten to post-fusion & only run once Triton was chosen for these 2 gemms so that's what I did. PiperOrigin-RevId: 974560034 --- .../xla/xla/backends/gpu/codegen/triton/BUILD | 5 +- .../codegen/triton/triton_gemm_fusion_test.cc | 270 +++++++++--------- 2 files changed, 142 insertions(+), 133 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/BUILD b/third_party/xla/xla/backends/gpu/codegen/triton/BUILD index 34cf8bfd3ca081..ec3ee6dfe836c5 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/triton/BUILD @@ -366,7 +366,6 @@ xla_test( "//xla/backends/gpu/tests:gpu_pjrt_codegen_test", "//xla/backends/gpu/transforms:convert_triton_gemm_config", "//xla/codegen/xtile:block_level_parameters", - "//xla/hlo/analysis:symbolic_map", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:filecheck", "//xla/hlo/testlib:pattern_matcher_gmock", @@ -379,8 +378,10 @@ xla_test( "//xla/service/gpu:target_constants", "//xla/stream_executor:device_description", "//xla/stream_executor/cuda:cuda_compute_capability", - "//xla/tests:hlo_pjrt_interpreter_reference_mixin", + "//xla/tests:hlo_interpreter_reference_mixin", "//xla/tests:xla_internal_test_main", # fixdeps: keep + "//xla/tsl/lib/core:status_test_util", + "//xla/tsl/platform:env", "//xla/tsl/platform:errors", "//xla/tsl/platform:test", "@com_google_absl//absl/log:check", diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc index 5ffa216b50edda..3186968ae3c3b3 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc @@ -40,7 +40,6 @@ limitations under the License. #include "xla/backends/gpu/transforms/convert_triton_gemm_config.h" #include "xla/codegen/xtile/block_level_parameters.h" #include "xla/error_spec.h" -#include "xla/hlo/analysis/symbolic_expr.h" #include "xla/hlo/ir/hlo_casting_utils.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" @@ -57,7 +56,9 @@ limitations under the License. #include "xla/service/pattern_matcher.h" #include "xla/stream_executor/cuda/cuda_compute_capability.h" #include "xla/stream_executor/device_description.h" -#include "xla/tests/hlo_pjrt_interpreter_reference_mixin.h" +#include "xla/tests/hlo_interpreter_reference_mixin.h" +#include "xla/tsl/lib/core/status_test_util.h" +#include "xla/tsl/platform/env.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/test.h" #include "xla/xla.pb.h" @@ -87,9 +88,9 @@ HloInstruction* GetNonBitcastRoot(const HloComputation* computation) { return root; } -class TritonTest : public HloInterpreterReferenceMixin { +class TritonTestBase : public HloInterpreterReferenceMixin { public: - TritonTest() = default; + TritonTestBase() = default; DebugOptions GetDebugOptionsForTest() const override { DebugOptions debug_options = GpuPjRtCodegenTest::GetDebugOptionsForTest(); debug_options.set_xla_gpu_autotune_level(0); @@ -134,6 +135,16 @@ class TritonTest : public HloInterpreterReferenceMixin { mlir::MLIRContext mlir_context_; }; +class TritonTest : public TritonTestBase, + public ::testing::WithParamInterface { + public: + DebugOptions GetDebugOptionsForTest() const override { + DebugOptions debug_options = TritonTestBase::GetDebugOptionsForTest(); + debug_options.set_xla_gpu_experimental_gemm_fusion_v2(GetParam()); + return debug_options; + } +}; + class TritonGemmTest : public TritonTest { public: DebugOptions GetDebugOptionsForTest() const override { @@ -158,7 +169,16 @@ class TritonGemmTest : public TritonTest { } }; -TEST_F(TritonTest, TestGemmWithTrivialNonContractingDimension) { +std::string GemmFusionVersionName(const ::testing::TestParamInfo& info) { + return info.param ? "GemmFusionV2" : "GemmFusionV1"; +} + +INSTANTIATE_TEST_SUITE_P(All, TritonTest, ::testing::Bool(), + GemmFusionVersionName); +INSTANTIATE_TEST_SUITE_P(All, TritonGemmTest, ::testing::Bool(), + GemmFusionVersionName); + +TEST_P(TritonTest, TestGemmWithTrivialNonContractingDimension) { constexpr absl::string_view kHloText = R"( HloModule t, is_scheduled=true @@ -189,7 +209,7 @@ CHECK: tt.dot {{.*}} : tensor<16x32xf32> * tensor<32x16xf32> -> tensor<16x16xf32 )")); } -TEST_F(TritonTest, PredParametersAreTruncatedToI1) { +TEST_P(TritonTest, PredParametersAreTruncatedToI1) { constexpr absl::string_view kHloText = R"( HloModule m @@ -233,7 +253,7 @@ CHECK: %{{.*}} = arith.andi %[[CMPI]], %{{.*}} : tensor<16x16xi1> )")); } -TEST_F(TritonTest, CodegenBatchedDotWithConcatenationWithCorrectBatchStride) { +TEST_P(TritonTest, CodegenBatchedDotWithConcatenationWithCorrectBatchStride) { constexpr absl::string_view kHloText = R"( HloModule t @@ -270,7 +290,7 @@ CHECK: tt.dot {{.*}} : tensor<16x32xf32> * tensor<32x64xf32> -> tensor<16x64xf32 } // TODO(b/417172838): enable after enabling dynamic slice in support.cc. -TEST_F(TritonTest, DISABLED_CodegenDynamicSliceWithCorrectOffsets) { +TEST_P(TritonTest, DISABLED_CodegenDynamicSliceWithCorrectOffsets) { // TODO(b/417172838): we now should support non-majormost dimensions, port // this test to fusion_emitter_device_test with that support. @@ -325,7 +345,7 @@ ENTRY e { )")); } -TEST_F(TritonGemmTest, DoNotUseTensorCoresWithHighestPrecision) { +TEST_P(TritonGemmTest, DoNotUseTensorCoresWithHighestPrecision) { constexpr absl::string_view kHloText = R"( triton_gemm_r { parameter_0 = s8[80,15]{1,0} parameter(0) @@ -355,7 +375,7 @@ CHECK-NOT: mma )"); } -TEST_F(TritonGemmTest, UseTensorCoresWithHighPrecision) { +TEST_P(TritonGemmTest, UseTensorCoresWithHighPrecision) { constexpr absl::string_view kHloText = R"( triton_gemm_r { parameter_0 = s8[80,15]{1,0} parameter(0) @@ -385,7 +405,7 @@ CHECK: mma )"); } -TEST_F(TritonGemmTest, DebugOptionsArePropagated) { +TEST_P(TritonGemmTest, DebugOptionsArePropagated) { constexpr absl::string_view kHloText = R"( ENTRY e { p0 = f16[30,30] parameter(0) @@ -423,7 +443,7 @@ ENTRY e { EXPECT_GT(file_size, 10); } -TEST_F(TritonGemmTest, DotWithPredFromCompareProducesCorrectResult) { +TEST_P(TritonGemmTest, DotWithPredFromCompareProducesCorrectResult) { constexpr absl::string_view kHloText = R"( triton_dot { parameter_0 = s32[4,128]{1,0} parameter(0) @@ -449,7 +469,7 @@ ENTRY main { ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-6})); } -TEST_F(TritonGemmTest, UseTensorCoresForF32OnAmpere) { +TEST_P(TritonGemmTest, UseTensorCoresForF32OnAmpere) { constexpr absl::string_view kHloText = R"( triton_gemm_r { parameter_0 = f16[80,15]{1,0} parameter(0) @@ -479,7 +499,7 @@ CHECK: mma // TODO(b/353484968): Tests that don't run RunAndCompareNoHloPasses should be // moved to deviceless test file. -TEST_F(TritonGemmTest, FailIfTooMuchShmem) { +TEST_P(TritonGemmTest, FailIfTooMuchShmem) { auto cc = se::CudaComputeCapability::Ampere(); const se::DeviceDescription device_info = TestGpuDeviceInfo::RTXA6000DeviceInfo(); @@ -536,7 +556,7 @@ ENTRY entry { EXPECT_GT(result.shmem_bytes, device_info.shared_memory_per_block()); } -TEST_F(TritonGemmTest, MultipleDims) { +TEST_P(TritonGemmTest, MultipleDims) { constexpr absl::string_view kHloText = R"( HloModule t @@ -559,7 +579,7 @@ ENTRY e { EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); } -TEST_F(TritonGemmTest, MultipleBatchDimensions) { +TEST_P(TritonGemmTest, MultipleBatchDimensions) { constexpr absl::string_view kHloText = R"( HloModule m @@ -593,7 +613,7 @@ ENTRY e { std::move(module), ErrorSpec{/*aabs=*/2e-2, /*arel=*/2e-2})); } -TEST_F(TritonGemmTest, PredWithBF16DotProducesCorrectResult) { +TEST_P(TritonGemmTest, PredWithBF16DotProducesCorrectResult) { constexpr absl::string_view kHloText = R"( triton_dot { p0 = pred[8,640]{1,0} parameter(0) @@ -622,7 +642,7 @@ ENTRY e { // TODO: b/422676780 - Enable the tests once the indexing maps-based tiling is // deprecated. The test is disabled after we remove TransposeDimensionGrouper // pass, because the infra currently requires grouping of adjacent dimensions. -TEST_F(TritonGemmTest, DISABLED_SplitLhsNoncontractingTransposeRhs) { +TEST_P(TritonGemmTest, DISABLED_SplitLhsNoncontractingTransposeRhs) { constexpr absl::string_view kHloText = R"( HloModule t @@ -650,7 +670,7 @@ ENTRY e { // TODO: b/422676780 - Enable the tests once the indexing maps-based tiling is // deprecated. The test is disabled after we remove TransposeDimensionGrouper // pass, because the infra currently requires grouping of adjacent dimensions. -TEST_F(TritonGemmTest, DISABLED_SplitLhsNoncontracting) { +TEST_P(TritonGemmTest, DISABLED_SplitLhsNoncontracting) { constexpr absl::string_view kHloText = R"( ENTRY e { p0 = f32[72,72] parameter(0) @@ -677,7 +697,7 @@ ENTRY e { EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); } -TEST_F(TritonGemmTest, SplitAndTransposeLhsExecutesCorrectly) { +TEST_P(TritonGemmTest, SplitAndTransposeLhsExecutesCorrectly) { constexpr absl::string_view kHloText = R"( HloModule m @@ -710,7 +730,7 @@ ENTRY e { // TODO(b/393299275): it's not clear that this test is actually testing what it // claims to be testing. It should either be rewritten to start from // post-optimization HLO, or hoisted out to test the fusion logic specifically. -TEST_F(TritonGemmTest, NondefaultOperandLayoutIsSupported) { +TEST_P(TritonGemmTest, NondefaultOperandLayoutIsSupported) { constexpr absl::string_view kHloText = R"( ENTRY r { p1 = f16[3,10,128]{2,1,0} parameter(1) @@ -723,10 +743,10 @@ ENTRY r { })"; MatchOptimizedHlo(kHloText, R"( -; CHECK: %[[p0:.*]] = f16[10,3,128]{2,0,1} parameter(0) +; CHECK: %[[p0:.*]] = f16[10,3,128]{2,0,1} parameter ; CHECK: %[[cv:.*]] = f32[10,3,128]{2,0,1} convert(%[[p0]]) ; CHECK: %[[tr:.*]] = f32[3,10,128]{2,0,1} transpose(%[[cv]]), dimensions={1,0,2} -; CHECK: %[[p1:.*]] = f32[3,10,123]{2,1,0} parameter(1) +; CHECK: %[[p1:.*]] = f32[3,10,123]{2,1,0} parameter ; CHECK: f32[3,128,123]{2,1,0} dot(%[[tr]], %[[p1]]) )"); @@ -735,7 +755,7 @@ ENTRY r { // TODO(b/393299275): this is a pure test of fusion logic. It should be moved to // a separate, fusion-specific, deviceless test. -TEST_F(TritonGemmTest, DoNotFuseSplitRhsContractingTranspose) { +TEST_P(TritonGemmTest, DoNotFuseSplitRhsContractingTranspose) { constexpr absl::string_view kHloText = R"( HloModule t @@ -761,7 +781,7 @@ ENTRY e { // TODO(b/393299275): this is a pure test of fusion logic. It should be moved to // a separate, fusion-specific, deviceless test. -TEST_F(TritonGemmTest, DoNotFuseSplitLhsContractingTranspose) { +TEST_P(TritonGemmTest, DoNotFuseSplitLhsContractingTranspose) { constexpr absl::string_view kHloText = R"( HloModule t @@ -788,7 +808,7 @@ ENTRY e { // TODO(b/393299275): this test should be rewritten to start from // post-optimization HLO. (Though I'm not entirely sure it's even worth keeping // it.) -TEST_F(TritonGemmTest, NonMajorMostInputBatchWorksCorrectly) { +TEST_P(TritonGemmTest, NonMajorMostInputBatchWorksCorrectly) { constexpr absl::string_view kHloText = R"( HloModule t @@ -815,7 +835,7 @@ ENTRY e { // TODO(b/393299275): this is a pure test of fusion logic. It should be moved to // a separate, fusion-specific, deviceless test. -TEST_F(TritonGemmTest, DoNotFuseArbitraryReshape) { +TEST_P(TritonGemmTest, DoNotFuseArbitraryReshape) { constexpr absl::string_view kHloText = R"( HloModule m @@ -838,36 +858,12 @@ ENTRY e { )"); } -TEST_F(TritonGemmTest, MultipleBatchRequireSeparateTranspose) { - constexpr absl::string_view kHloText = R"( -HloModule m - -ENTRY e { - Arg_0 = f16[3,4,2,5,4] parameter(0) - c = f32[3,4,2,5,4] convert(Arg_0) - Arg_1 = f32[5,3,4,3,2] parameter(1) - ROOT dot.3 = f32[5,3,4,4,3] dot(c, Arg_1), - lhs_batch_dims={3,0,1}, lhs_contracting_dims={2}, - rhs_batch_dims={0,1,2}, rhs_contracting_dims={4} -})"; - - MatchOptimizedHlo(kHloText, R"( -; CHECK: ROOT -; CHECK: transpose( -; CHECK: bitcast( -; CHECK: kCustom -; CHECK-SAME: backend_config={{.*}}"kind":"__triton_nested_gemm_fusion" -)"); - - EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-4})); -} - // This tests the complexity heuristics in TritonWrapper. // TODO(b/393299275): this is not worth keeping as a codegen test. Really, we // should not reject tilings that are slow/spill in codegen. If this has use in // autotuning, then this should be tested/called in the autotuner. // The generic Triton emitter does not want to deal with this. -TEST_F(TritonGemmTest, DISABLED_FailForTooComplexTiling) { +TEST_P(TritonGemmTest, DISABLED_FailForTooComplexTiling) { auto cc = se::CudaComputeCapability::Ampere(); const se::DeviceDescription device_info = TestGpuDeviceInfo::RTXA6000DeviceInfo(); @@ -928,7 +924,7 @@ ENTRY entry { // TODO(b/393299275): this test may have some value while Triton tiling // propagation is being replaced, but has little worth as a codegen test. // Consider moving this. -TEST_F(TritonGemmTest, +TEST_P(TritonGemmTest, BroadcastsOfTriviallySizedNonContractingDimensionsAreSupported) { constexpr absl::string_view kHloText = R"( f { @@ -963,7 +959,7 @@ e { // TODO(b/393299275): this test may have some value while Triton tiling // propagation is being replaced, but has little worth as a codegen test. // Consider moving this. -TEST_F(TritonGemmTest, +TEST_P(TritonGemmTest, BroadcastsOfTriviallySizedContractingDimensionsAreSupported) { constexpr absl::string_view kHloText = R"( f { @@ -994,7 +990,7 @@ e { } // TODO(b/417172838): enable after enabling dynamic slice in support.cc. -TEST_F(TritonGemmTest, DISABLED_DynamicSliceIsSupportedInLhsEndToEnd) { +TEST_P(TritonGemmTest, DISABLED_DynamicSliceIsSupportedInLhsEndToEnd) { // The select is used to restrict the start index to values that make sense. // If it was constant, then the dynamic-slice would be optimized to slice. It // is not strictly needed, because we also support clamping the indices. @@ -1029,7 +1025,7 @@ ENTRY e { } // TODO(b/417172838): enable after enabling dynamic slice in support.cc. -TEST_F(TritonGemmTest, DISABLED_DynamicSliceIsSupportedInRhs) { +TEST_P(TritonGemmTest, DISABLED_DynamicSliceIsSupportedInRhs) { // The start index(es) for the non-majormost dimension(s) are constant zero(s) // because we don't support dynamic slice on those dimensions. constexpr absl::string_view kHloText = R"( @@ -1067,7 +1063,7 @@ ENTRY e { } class TritonGemmDynamicSliceClampingTest - : public TritonTest, + : public TritonTestBase, public ::testing::WithParamInterface {}; // TODO(b/417172838): enable after enabling dynamic slice in support.cc. @@ -1122,7 +1118,7 @@ INSTANTIATE_TEST_SUITE_P(All, TritonGemmDynamicSliceClampingTest, ::testing::Values(-100, 3, 999), OffsetParamToString); // TODO(b/417172838): enable after enabling dynamic slice in support.cc. -TEST_F(TritonGemmTest, +TEST_P(TritonGemmTest, DISABLED_DynamicSliceOfMajormostContractingDimIsSupported) { // Tests that dynamic-slice works on the majormost dimension even if that // dimension is contracted. @@ -1163,7 +1159,7 @@ ENTRY e { } // TODO(b/417172838): enable after enabling dynamic slice in support.cc. -TEST_F(TritonGemmTest, DISABLED_DynamicSliceOfMajormostBatchDimIsSupported) { +TEST_P(TritonGemmTest, DISABLED_DynamicSliceOfMajormostBatchDimIsSupported) { // Tests that dynamic-slice works on the majormost dimension even if that // dimension is a batch. // The start index(es) for the non-majormost dimension(s) are constant zero(s) @@ -1208,7 +1204,7 @@ ENTRY e { } // TODO(b/417172838): enable after enabling dynamic slice in support.cc. -TEST_F(TritonGemmTest, +TEST_P(TritonGemmTest, DISABLED_DynamicSliceSingleDimensionIntoReshapeIsSupported) { // This directly tests the targeted use case (b/307922364) of iterating over // layer weights and extracting them with dynamic slice. @@ -1253,7 +1249,7 @@ ENTRY e { // TODO(b/393299275): this should just be a fusion test and does not need to be // in the codegen directory. -TEST_F(TritonGemmTest, DoNotFuseConcatenationOfSplitNonContractingDimension) { +TEST_P(TritonGemmTest, DoNotFuseConcatenationOfSplitNonContractingDimension) { if (GpuComputeCapability().IsRocm()) { GTEST_SKIP() << "Not using autotuner on ROCM yet."; } @@ -1283,7 +1279,7 @@ ENTRY e { EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); } -TEST_F(TritonGemmTest, BinaryOperationWithSmallInputsIsFused) { +TEST_P(TritonGemmTest, BinaryOperationWithSmallInputsIsFused) { constexpr absl::string_view kHloText = R"( HloModule m @@ -1307,7 +1303,7 @@ ENTRY e { .WithFusionKind(HloInstruction::FusionKind::kCustom))); } -TEST_F(TritonGemmTest, BinaryOperationWithLargeInputsIsNotFused) { +TEST_P(TritonGemmTest, BinaryOperationWithLargeInputsIsNotFused) { constexpr absl::string_view kHloText = R"( HloModule m @@ -1336,7 +1332,7 @@ ENTRY e { )"); } -TEST_F(TritonGemmTest, ParametersWithDifferentLayoutsAreSupportedInOneScope) { +TEST_P(TritonGemmTest, ParametersWithDifferentLayoutsAreSupportedInOneScope) { constexpr absl::string_view kHloText = R"( ENTRY e { p0 = s8[5,3] parameter(0) @@ -1357,7 +1353,7 @@ ENTRY e { .WithFusionKind(HloInstruction::FusionKind::kCustom))); } -TEST_F(TritonGemmTest, BinaryOperationOnLargeParametersIsFused) { +TEST_P(TritonGemmTest, BinaryOperationOnLargeParametersIsFused) { constexpr absl::string_view kHloText = R"( HloModule m @@ -1380,45 +1376,55 @@ ENTRY e { .WithFusionKind(HloInstruction::FusionKind::kCustom))); } -TEST_F(TritonGemmTest, LinkingLibdeviceTwiceWorks) { - // TODO(b/393299275): This test looks weird. It's testing the whole - // optimization pipeline end-to-end to check that linking libdevice twice - // works? rewrite this to just use post-optimization HLO +TEST_P(TritonGemmTest, LinkingLibdeviceTwiceWorks) { constexpr absl::string_view kHloText = R"( -ENTRY e { - p0 = s8[7,3] parameter(0) - c0 = f32[7,3] convert(p0) - p1 = f32[3,16] parameter(1) - e1 = f32[3,16] exponential(p1) - d0 = f32[7,16] dot(c0, e1), - lhs_contracting_dims={1}, rhs_contracting_dims={0} - p2 = s8[7,3] parameter(2) - c2 = f32[7,3] convert(p2) - e2 = f32[7,3] exponential(c2) - p3 = f32[3,16] parameter(3) - d1 = f32[7,16] dot(e2, p3), +HloModule m + +triton_dot_0 { + p0 = s8[7,3]{1,0} parameter(0) + c0 = f32[7,3]{1,0} convert(p0) + p1 = f32[3,16]{1,0} parameter(1) + e1 = f32[3,16]{1,0} exponential(p1) + ROOT d0 = f32[7,16]{1,0} dot(c0, e1), lhs_contracting_dims={1}, rhs_contracting_dims={0} - ROOT a = f32[7,16] add(d0, d1) -})"; +} - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetOptimizedModule(kHloText)); +triton_dot_1 { + p0 = s8[7,3]{1,0} parameter(0) + c0 = f32[7,3]{1,0} convert(p0) + e0 = f32[7,3]{1,0} exponential(c0) + p1 = f32[3,16]{1,0} parameter(1) + ROOT d1 = f32[7,16]{1,0} dot(e0, p1), + lhs_contracting_dims={1}, rhs_contracting_dims={0} +} - EXPECT_THAT( - module->entry_computation()->root_instruction(), - GmockMatch( - m::Fusion(m::Fusion(m::Parameter(), m::Parameter()) - .WithFusionKind(HloInstruction::FusionKind::kCustom), - m::Fusion(m::Parameter(), m::Parameter()) - .WithFusionKind(HloInstruction::FusionKind::kCustom)) - .WithFusionKind(HloInstruction::FusionKind::kLoop))); +ENTRY e { + p0 = s8[7,3]{1,0} parameter(0) + p1 = f32[3,16]{1,0} parameter(1) + p2 = s8[7,3]{1,0} parameter(2) + p3 = f32[3,16]{1,0} parameter(3) + f0 = f32[7,16]{1,0} fusion(p0, p1), kind=kCustom, calls=triton_dot_0, + backend_config={"fusion_backend_config": {kind: "__triton_gemm", + triton_gemm_config: {"block_m":16,"block_n":16,"block_k":16,"num_stages":1,"num_warps":2,"num_ctas":1}}} + f1 = f32[7,16]{1,0} fusion(p2, p3), kind=kCustom, calls=triton_dot_1, + backend_config={"fusion_backend_config": {kind: "__triton_gemm", + triton_gemm_config: {"block_m":16,"block_n":16,"block_k":16,"num_stages":1,"num_warps":2,"num_ctas":1}}} + ROOT a = f32[7,16]{1,0} add(f0, f1) +})"; - EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-2, /*arel=*/1e-2})); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kHloText)); + ASSERT_OK_AND_ASSIGN(bool converted, ConvertTritonGemmConfig( + device_description(), &mlir_context_) + .Run(module.get())); + EXPECT_TRUE(converted); + EXPECT_TRUE(RunAndCompareNoHloPasses( + std::move(module), ErrorSpec{/*aabs=*/1e-2, /*arel=*/1e-2})); } // TODO(b/393299275): this should just be a fusion test and does not need to be // in the codegen directory. -TEST_F(TritonGemmTest, BroadcastOfScalarParameterIsFused) { +TEST_P(TritonGemmTest, BroadcastOfScalarParameterIsFused) { constexpr absl::string_view kHloText = R"( ENTRY e { p0 = f16[64,256] parameter(0) @@ -1439,7 +1445,7 @@ ENTRY e { // TODO(b/393299275): this should just be a fusion test and does not need to be // in the codegen directory. -TEST_F(TritonGemmTest, BroadcastOfScalarConstantIsFused) { +TEST_P(TritonGemmTest, BroadcastOfScalarConstantIsFused) { constexpr absl::string_view kHloText = R"( HloModule m @@ -1462,7 +1468,7 @@ ENTRY e { // TODO(b/393299275): this should just be a fusion test and does not need to be // in the codegen directory. -TEST_F(TritonGemmTest, DoubleBroadcastOfScalarConstantIsFused) { +TEST_P(TritonGemmTest, DoubleBroadcastOfScalarConstantIsFused) { constexpr absl::string_view kHloText = R"( ENTRY e { c = s32[] constant(1) @@ -1486,7 +1492,7 @@ ENTRY e { // TODO(b/393299275): this should just be a fusion test and does not need to be // in the codegen directory. -TEST_F(TritonGemmTest, BroadcastOfVectorConstantIsFused) { +TEST_P(TritonGemmTest, BroadcastOfVectorConstantIsFused) { constexpr absl::string_view kHloText = R"( HloModule m @@ -1508,7 +1514,7 @@ ENTRY e { // TODO(b/393299275): this should just be a fusion test and does not need to be // in the codegen directory. -TEST_F(TritonGemmTest, AlwaysFuseScalarConstantAtBroadcastInput) { +TEST_P(TritonGemmTest, AlwaysFuseScalarConstantAtBroadcastInput) { constexpr absl::string_view kHloText = R"( ENTRY e { p0 = bf16[2,3,3]{2,1,0} parameter(0) @@ -1536,7 +1542,7 @@ ENTRY e { // TODO(b/393299275): this should just be a fusion test and does not need to be // in the codegen directory. -TEST_F(TritonGemmTest, FuseConcatenation) { +TEST_P(TritonGemmTest, FuseConcatenation) { if (!SupportsBF16(GpuComputeCapability())) { GTEST_SKIP() << "BF16 not supported."; } @@ -1566,7 +1572,7 @@ e { // TODO(b/393299275): this should just be a fusion test and does not need to be // in the codegen directory. -TEST_F(TritonGemmTest, SineOutputIsFused) { +TEST_P(TritonGemmTest, SineOutputIsFused) { constexpr absl::string_view kHloText = R"( HloModule m @@ -1589,7 +1595,7 @@ ENTRY e { // TODO(b/393299275): this should just be a fusion test and does not need to be // in the codegen directory. -TEST_F(TritonGemmTest, SliceInputIsFused) { +TEST_P(TritonGemmTest, SliceInputIsFused) { constexpr absl::string_view kHloText = R"( ENTRY e { p0 = f16[97,121] parameter(0) @@ -1610,7 +1616,7 @@ ENTRY e { // TODO(b/393299275): this should just be a fusion test and does not need to be // in the codegen directory. -TEST_F(TritonGemmTest, SliceInputWithReshapeIsFused) { +TEST_P(TritonGemmTest, SliceInputWithReshapeIsFused) { constexpr absl::string_view kHloText = R"( ENTRY e { p0 = f32[363,1536] parameter(0) @@ -1635,7 +1641,7 @@ ENTRY e { // TODO(b/393299275): this should just be a fusion test and does not need to be // in the codegen directory. From the point of view of codegen, this is purely // about compositionality. -TEST_F(TritonGemmTest, NestedSlicingWorks) { +TEST_P(TritonGemmTest, NestedSlicingWorks) { constexpr absl::string_view kHloText = R"( ENTRY e { p1 = f32[6,24] parameter(1) @@ -1657,7 +1663,7 @@ ENTRY e { // TODO(b/393299275): this should just be a fusion test and does not need to be // in the codegen directory. -TEST_F(TritonGemmTest, SlicedBatchDimensionIsSupported) { +TEST_P(TritonGemmTest, SlicedBatchDimensionIsSupported) { constexpr absl::string_view kHloText = R"( ENTRY e { p0 = f16[3,3,256] parameter(0) @@ -1681,7 +1687,7 @@ ENTRY e { EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); } -TEST_F(TritonGemmTest, BatchMajorSlicedBatchDimensionProducesCorrectResults) { +TEST_P(TritonGemmTest, BatchMajorSlicedBatchDimensionProducesCorrectResults) { constexpr absl::string_view kHloText = R"( ENTRY e { p0 = f16[4,32,256] parameter(0) @@ -1694,7 +1700,7 @@ ENTRY e { EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); } -TEST_F(TritonGemmTest, DegenerateBatchDimensionProducesCorrectResults) { +TEST_P(TritonGemmTest, DegenerateBatchDimensionProducesCorrectResults) { constexpr absl::string_view kHloText = R"( ENTRY e { p0 = f16[1,32,256] parameter(0) @@ -1709,7 +1715,7 @@ ENTRY e { // TODO(b/393299275): this should just be a fusion test and does not need to be // in the codegen directory. -TEST_F(TritonGemmTest, NarrowingConvertOutputIsFused) { +TEST_P(TritonGemmTest, NarrowingConvertOutputIsFused) { constexpr absl::string_view kHloText = R"( HloModule m @@ -1732,7 +1738,7 @@ ENTRY e { // TODO(b/393299275): this should just be a fusion test and does not need to be // in the codegen directory. -TEST_F(TritonGemmTest, ParameterAfterDotIsFused) { +TEST_P(TritonGemmTest, ParameterAfterDotIsFused) { if (!SupportsBF16(GpuComputeCapability())) { GTEST_SKIP() << "BF16 not supported."; } @@ -1764,7 +1770,7 @@ ENTRY e { EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/2e-2, /*arel=*/2e-2})); } -TEST_F(TritonGemmTest, OutputFusionExecutesCorrectly) { +TEST_P(TritonGemmTest, OutputFusionExecutesCorrectly) { // TODO(b/393299275): it this test useful? if (!SupportsBF16(GpuComputeCapability())) { GTEST_SKIP() << "BF16 not supported."; @@ -1803,7 +1809,7 @@ ENTRY e { // TODO(b/393299275): this should just be a fusion test and does not need to be // in the codegen directory. -TEST_F(TritonGemmTest, SplitLHSOutputTransposeAloneIsNotFused) { +TEST_P(TritonGemmTest, SplitLHSOutputTransposeAloneIsNotFused) { if (!SupportsBF16(GpuComputeCapability())) { GTEST_SKIP() << "BF16 not supported."; } @@ -1843,7 +1849,7 @@ ENTRY e { // time, the `bitcast` we need to hoist needs to be hoisted *upwards* but is // located after the dot (it collapses two consecutive non-contracting // dimensions together). -TEST_F(TritonGemmTest, DISABLED_SplitLHSInputOutputIsFused) { +TEST_P(TritonGemmTest, DISABLED_SplitLHSInputOutputIsFused) { if (!SupportsBF16(GpuComputeCapability())) { GTEST_SKIP() << "BF16 not supported."; } @@ -1874,7 +1880,7 @@ ENTRY e { EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); } -TEST_F(TritonGemmTest, SupportPredParametersUsedInExpressions) { +TEST_P(TritonGemmTest, SupportPredParametersUsedInExpressions) { constexpr absl::string_view kHloText = R"( ENTRY e { p = pred[2,2]{1,0} parameter(0) @@ -1904,7 +1910,7 @@ ENTRY e { // Triton emitter. This should probably be a fusion test. // TODO(b/353484968): Tests that don't run RunAndCompareNoHloPasses should be // moved to deviceless test file. -TEST_F(TritonGemmTest, LowerDotWithLhsWithoutNonContractingDimThroughTriton) { +TEST_P(TritonGemmTest, LowerDotWithLhsWithoutNonContractingDimThroughTriton) { constexpr absl::string_view kHloText = R"( ENTRY e { parameter_0 = f32[1,40] parameter(0) @@ -1927,7 +1933,7 @@ ENTRY e { // Triton emitter. This should probably be a fusion test. // TODO(b/353484968): Tests that don't run RunAndCompareNoHloPasses should be // moved to deviceless test file. -TEST_F(TritonGemmTest, LowerDotWithRhsWithoutNonContractingDimThroughTriton) { +TEST_P(TritonGemmTest, LowerDotWithRhsWithoutNonContractingDimThroughTriton) { constexpr absl::string_view kHloText = R"( ENTRY e { parameter_0 = f32[1,40,250000] parameter(0) @@ -1948,8 +1954,10 @@ ENTRY e { // This group of tests compares GPU results of dots already rewritten // into Triton fusions. using CompareTest = TritonGemmTest; +INSTANTIATE_TEST_SUITE_P(All, CompareTest, ::testing::Bool(), + GemmFusionVersionName); -TEST_F(CompareTest, F32WithTrivialNonContractingDimension) { +TEST_P(CompareTest, F32WithTrivialNonContractingDimension) { constexpr absl::string_view kHloTextRef = R"( HloModule r @@ -2001,7 +2009,7 @@ ENTRY e { // runs across every possible chip that is supported by Triton. The test should // probably be made deviceless and repurposed to test that opt-in shared memory // is used only. -TEST_F(CompareTest, UsingOptinSharedMemoryProducesSameResult) { +TEST_P(CompareTest, UsingOptinSharedMemoryProducesSameResult) { if (GpuComputeCapability().IsRocm()) { GTEST_SKIP() << "No Optin Shared Memory on AMD."; } @@ -2099,7 +2107,7 @@ ENTRY e { /*run_hlo_passes=*/false)); } -TEST_F(CompareTest, NonMajorMostOutputBatchWorksCorrectly) { +TEST_P(CompareTest, NonMajorMostOutputBatchWorksCorrectly) { const std::string kHloTextTest = R"( HloModule m @@ -2167,7 +2175,7 @@ ENTRY e { /*run_hlo_passes=*/false)); } -TEST_F(CompareTest, TritonDotFusionCanHaveOnlyRHSParameter) { +TEST_P(CompareTest, TritonDotFusionCanHaveOnlyRHSParameter) { const std::string kHloTextTest = R"( HloModule m, is_scheduled=true @@ -2215,7 +2223,7 @@ ENTRY e { /*run_hlo_passes=*/false)); } -TEST_F(CompareTest, TritonDotFusionCanHaveNoParametersAtAll) { +TEST_P(CompareTest, TritonDotFusionCanHaveNoParametersAtAll) { const std::string kHloTextTest = R"( HloModule m, is_scheduled=true @@ -2264,7 +2272,7 @@ ENTRY triton_gemm___computation { /*run_hlo_passes=*/false)); } -TEST_F(CompareTest, TritonDotFusionCanHaveManyParameters) { +TEST_P(CompareTest, TritonDotFusionCanHaveManyParameters) { const std::string kHloTextTest = R"( HloModule m @@ -2388,7 +2396,7 @@ ENTRY e { /*run_hlo_passes=*/false)); } -TEST_F(CompareTest, PredToBF16ConversionWorks) { +TEST_P(CompareTest, PredToBF16ConversionWorks) { if (!SupportsBF16(GpuComputeCapability())) { GTEST_SKIP() << "BF16 not supported."; } @@ -2459,7 +2467,7 @@ ENTRY e { // the two transposes in the fusion allows hoisting the final bitcast (`b1`). // I'm not sure if this is even required, since now we canonicalize transposes // before fusing. -TEST_F(CompareTest, DISABLED_DifferentLayoutsAreSupportedInOneScope) { +TEST_P(CompareTest, DISABLED_DifferentLayoutsAreSupportedInOneScope) { const std::string kHloTextTest = R"( triton_dot { p1 = f16[3,3,2,16]{1,3,2,0} parameter(1) @@ -2524,7 +2532,7 @@ ENTRY e { /*run_hlo_passes=*/false)); } -TEST_F(TritonTest, UseTF32For8BitOrLessWithF32) { +TEST_P(TritonTest, UseTF32For8BitOrLessWithF32) { constexpr absl::string_view kHloText = R"( HloModule t @@ -2565,7 +2573,7 @@ CHECK: inputPrecision = tf32 ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); } -TEST_F(TritonGemmTest, S8ToF16DotWithSmallTileDoesNotCrash) { +TEST_P(TritonGemmTest, S8ToF16DotWithSmallTileDoesNotCrash) { constexpr absl::string_view kHloText = R"( HloModule m @@ -2593,7 +2601,7 @@ ENTRY e { /*run_hlo_passes=*/false)); } -TEST_F(TritonGemmTest, S8ToF32DotWithManyWarpsDoesNotCrash) { +TEST_P(TritonGemmTest, S8ToF32DotWithManyWarpsDoesNotCrash) { constexpr absl::string_view kHloText = R"( HloModule m @@ -2621,7 +2629,7 @@ ENTRY e { /*run_hlo_passes=*/false)); } -TEST_F(TritonGemmTest, Fp8DotWithSmallTileDoesNotCrash) { +TEST_P(TritonGemmTest, Fp8DotWithSmallTileDoesNotCrash) { if (!GetCudaComputeCapability().IsAtLeastHopper()) { GTEST_SKIP() << "Doesn't pass on pre-Hopper GPUs."; } @@ -2652,7 +2660,7 @@ ENTRY e { /*run_hlo_passes=*/false)); } -TEST_F(TritonGemmTest, MixedF8DotExecutesCorrectly) { +TEST_P(TritonGemmTest, MixedF8DotExecutesCorrectly) { if (!GetCudaComputeCapability().IsAtLeastHopper()) { GTEST_SKIP() << "Requires a Hopper+ GPU"; } @@ -2693,7 +2701,7 @@ e { /*run_hlo_passes=*/false)); } -TEST_F(TritonGemmTest, Fp8DotWithManyWarpsDoesNotCrash) { +TEST_P(TritonGemmTest, Fp8DotWithManyWarpsDoesNotCrash) { if (!GetCudaComputeCapability().IsAtLeastHopper()) { GTEST_SKIP() << "Doesn't pass on pre-Hopper GPUs."; } @@ -2726,7 +2734,7 @@ ENTRY e { // TODO(b/353484968): Tests that don't run RunAndCompareNoHloPasses should be // moved to deviceless test file. -TEST_F(TritonGemmTest, WgmmaIsUsedForMemBoundShape) { +TEST_P(TritonGemmTest, WgmmaIsUsedForMemBoundShape) { if (GetCudaComputeCapability().major != se::CudaComputeCapability::kHopper) { GTEST_SKIP() << "wgmma instruction is only available on Hopper"; } @@ -2764,7 +2772,7 @@ CHECK: wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16 // // TODO(b/353484968): Tests that don't run RunAndCompareNoHloPasses should be // moved to deviceless test file. -TEST_F(TritonGemmTest, TestNoAutotuner) { +TEST_P(TritonGemmTest, TestNoAutotuner) { constexpr absl::string_view kHloText = R"( ENTRY e { p0 = f16[30,30] parameter(0) From 93066fa409f0683cad745e9faee095e05ee26a7b Mon Sep 17 00:00:00 2001 From: Bhatu Date: Tue, 1 Sep 2026 10:29:32 -0700 Subject: [PATCH 13/16] Propagate constraints across kCall boundaries in ConstraintPropagator. PiperOrigin-RevId: 974572981 --- .../xla/xla/tests/constraint_propagator.cc | 54 +++++++++++-------- .../xla/xla/tests/constraint_propagator.h | 11 ++-- .../xla/tests/constraint_propagator_test.cc | 29 ++++++++++ 3 files changed, 68 insertions(+), 26 deletions(-) diff --git a/third_party/xla/xla/tests/constraint_propagator.cc b/third_party/xla/xla/tests/constraint_propagator.cc index 15399c6d887e8f..70b33e7bd1000a 100644 --- a/third_party/xla/xla/tests/constraint_propagator.cc +++ b/third_party/xla/xla/tests/constraint_propagator.cc @@ -706,6 +706,15 @@ absl::Status ConstraintPropagator::SeedConstraints( states_[fused_computation->parameter_instruction(i)]; states_[inst->operand(i)].Merge(source_state); } + } else if (inst->opcode() == HloOpcode::kCall) { + const HloComputation* called_computation = inst->to_apply(); + if (called_computation != nullptr) { + for (int i = 0; i < inst->operand_count(); ++i) { + ConstraintState source_state = + states_[called_computation->parameter_instruction(i)]; + states_[inst->operand(i)].Merge(source_state); + } + } } } return absl::OkStatus(); @@ -865,7 +874,11 @@ absl::Status ConstraintPropagator::SeedMLPatternsConstraints( absl::Status ConstraintPropagator::PropagateConstraintsExact( const HloInstruction* instruction) { if (instruction->opcode() == HloOpcode::kFusion) { - return PropagateFusionBoundary(instruction); + return PropagateComputationBoundary( + instruction, instruction->fused_instructions_computation()); + } + if (instruction->opcode() == HloOpcode::kCall) { + return PropagateComputationBoundary(instruction, instruction->to_apply()); } ConstraintState output_state = states_[instruction]; @@ -989,41 +1002,40 @@ absl::Status ConstraintPropagator::PropagateConstraintsExact( return absl::OkStatus(); } -absl::Status ConstraintPropagator::PropagateFusionBoundary( - const HloInstruction* fusion_instruction) { - const HloComputation* fused_comp = - fusion_instruction->fused_instructions_computation(); - if (fused_comp == nullptr) { +absl::Status ConstraintPropagator::PropagateComputationBoundary( + const HloInstruction* caller_instruction, + const HloComputation* callee_computation) { + if (callee_computation == nullptr) { return absl::OkStatus(); } // 1. Output / Root binding: - const HloInstruction* fused_root = - fusion_instruction->fused_expression_root(); - if (fused_root != nullptr) { - ConstraintState fusion_state = states_[fusion_instruction]; - ConstraintState root_state = states_[fused_root]; - // Backward: outer constraint on fusion result flows into inner root. - states_[fused_root].Merge(fusion_state); - // Forward: internal constraint computed on root flows out to fusion result. - states_[fusion_instruction].Merge(root_state); + const HloInstruction* callee_root = callee_computation->root_instruction(); + if (callee_root != nullptr) { + ConstraintState caller_state = states_[caller_instruction]; + ConstraintState root_state = states_[callee_root]; + // Backward: outer constraint on caller result flows into inner root. + states_[callee_root].Merge(caller_state); + // Forward: internal constraint computed on root flows out to caller result. + states_[caller_instruction].Merge(root_state); } // 2. Operands / Parameters binding: - for (int64_t i = 0; i < fusion_instruction->operand_count(); ++i) { - const HloInstruction* operand = fusion_instruction->operand(i); - const HloInstruction* fused_param = fusion_instruction->fused_parameter(i); - if (fused_param == nullptr) { + for (int64_t i = 0; i < caller_instruction->operand_count(); ++i) { + const HloInstruction* operand = caller_instruction->operand(i); + const HloInstruction* callee_param = + callee_computation->parameter_instruction(i); + if (callee_param == nullptr) { continue; } ConstraintState operand_state = states_[operand]; - ConstraintState param_state = states_[fused_param]; + ConstraintState param_state = states_[callee_param]; // Backward: constraints accumulated on the internal parameter flow out // to the caller operand. states_[operand].Merge(param_state); // Forward: constraints established on the caller operand flow into the // internal parameter. - states_[fused_param].Merge(operand_state); + states_[callee_param].Merge(operand_state); } return absl::OkStatus(); diff --git a/third_party/xla/xla/tests/constraint_propagator.h b/third_party/xla/xla/tests/constraint_propagator.h index d70045c768df57..f78f34369598c5 100644 --- a/third_party/xla/xla/tests/constraint_propagator.h +++ b/third_party/xla/xla/tests/constraint_propagator.h @@ -108,11 +108,12 @@ class ConstraintPropagator { // formatting can simply propagate the exact constraints to their operands. absl::Status PropagateConstraintsExact(const HloInstruction* instruction); - // Propagates constraints bidirectionally across a kFusion instruction - // boundary, mapping caller operands to callee fused parameters and caller - // result to the fused expression root. - absl::Status PropagateFusionBoundary( - const HloInstruction* fusion_instruction); + // Propagates constraints bidirectionally across a subcomputation boundary + // (e.g. kFusion or kCall), mapping caller operands to callee parameters and + // caller result to the callee root instruction. + absl::Status PropagateComputationBoundary( + const HloInstruction* caller_instruction, + const HloComputation* callee_computation); // Propagates constraints from the output of an instruction to its operands. // This is approximate and introduces approximations for ops like add, sub, diff --git a/third_party/xla/xla/tests/constraint_propagator_test.cc b/third_party/xla/xla/tests/constraint_propagator_test.cc index 340d0f881b65aa..04c53c8656f156 100644 --- a/third_party/xla/xla/tests/constraint_propagator_test.cc +++ b/third_party/xla/xla/tests/constraint_propagator_test.cc @@ -1365,5 +1365,34 @@ ENTRY main { EXPECT_TRUE(p0_int.IsPositive()); EXPECT_GE(p0_int.min, 0.0); } + +TEST_F(ConstraintPropagatorTest, CallPropagates) { + const char* hlo = R"( +HloModule TestModule + +%sc_gather_comp (operand: f32[32,128], indices: s32[256]) -> f32[256,128] { + %operand = f32[32,128] parameter(0) + %indices = s32[256] parameter(1) + ROOT %gather = f32[256,128] gather(%operand, %indices), + offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, + index_vector_dim=1, slice_sizes={1,128} +} + +ENTRY main { + %operand = f32[32,128] parameter(0) + %indices = s32[256] parameter(1) + ROOT %call = f32[256,128] call(%operand, %indices), to_apply=%sc_gather_comp +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + ASSERT_OK_AND_ASSIGN(auto states, ConstraintPropagator::Run(*module)); + + auto indices_int = + states[module->entry_computation()->parameter_instruction(1)] + .GetConstraintInterval(); + EXPECT_FALSE(indices_int.IsEmpty()); + EXPECT_DOUBLE_EQ(indices_int.min, 0.0); + EXPECT_DOUBLE_EQ(indices_int.max, 31.0); +} } // namespace } // namespace xla From 26a3cc9cd3b576808399cc7d2fa3e9bd67e12423 Mon Sep 17 00:00:00 2001 From: Tori Baker Date: Tue, 1 Sep 2026 10:47:08 -0700 Subject: [PATCH 14/16] Get TAP ready for Gemm Fusion V2 flag flip - explicitly set experimental_gemm_fusion flag in each case that it's meant to represent. These will get cleaned up after the flag flip, but explicitly setting it means that the tests themselves are not affected (they run both versions when they're meant to or only 1 version where they differ intentionally). PiperOrigin-RevId: 974584295 --- .../gpu/transforms/gemm_fusion_test.cc | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc index c1c6e89ad2b2bc..7022507b428b44 100644 --- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc @@ -78,6 +78,7 @@ class GemmFusionTestBase : public HloHardwareIndependentTestBase { HloHardwareIndependentTestBase::GetDebugOptionsForTest(); debug_options.set_xla_gpu_triton_gemm_any(false); debug_options.set_xla_gpu_gemm_rewrite_size_threshold(0); + debug_options.set_xla_gpu_experimental_gemm_fusion_v2(false); return debug_options; } @@ -96,15 +97,16 @@ class GemmFusionTest : public GemmFusionTestBase, public: DebugOptions GetDebugOptionsForTest() const override { DebugOptions debug_options = GemmFusionTestBase::GetDebugOptionsForTest(); - if (GetParam()) { - debug_options.set_xla_gpu_experimental_enable_tiling_propagation(true); - } + debug_options.set_xla_gpu_experimental_enable_tiling_propagation( + GetParam()); return debug_options; } }; -// Create a parameterized test that makes sure that both the legacy and the new -// implementation of dot fusion are working as expected. +// While we launch, create a parameterized test to test all combinations. This +// test class has 2 parameters: +// 1. Whether to use Gemm Fusion V1 or V2. +// 2. Whether to use symbolic analysis or tiling propagation. class GemmFusionTestVersioned : public GemmFusionTestBase, public ::testing::WithParamInterface> { @@ -113,9 +115,8 @@ class GemmFusionTestVersioned DebugOptions debug_options = GemmFusionTestBase::GetDebugOptionsForTest(); debug_options.set_xla_gpu_experimental_gemm_fusion_v2( std::get<0>(GetParam())); - if (std::get<1>(GetParam())) { - debug_options.set_xla_gpu_experimental_enable_tiling_propagation(true); - } + debug_options.set_xla_gpu_experimental_enable_tiling_propagation( + std::get<1>(GetParam())); return debug_options; } }; @@ -128,9 +129,8 @@ class GemmFusionTestV2 : public GemmFusionTestBase, DebugOptions GetDebugOptionsForTest() const override { DebugOptions debug_options = GemmFusionTestBase::GetDebugOptionsForTest(); debug_options.set_xla_gpu_experimental_gemm_fusion_v2(true); - if (GetParam()) { - debug_options.set_xla_gpu_experimental_enable_tiling_propagation(true); - } + debug_options.set_xla_gpu_experimental_enable_tiling_propagation( + GetParam()); return debug_options; } }; From 2d7a460404404d2291e43cb56a5400730882cb49 Mon Sep 17 00:00:00 2001 From: Junwhan Ahn Date: Tue, 1 Sep 2026 10:47:09 -0700 Subject: [PATCH 15/16] Leverage input_devices_for_output_map in RemapArrays for PJRT IFRT. When an optimized remap plan provides `input_devices_for_output_map`, it explicitly details which input arrays and device shards contribute to each output array. Previously, `PjRtCompatibleClientRemapArrays` always inspected `plan.mappings()` regardless of whether `input_devices_for_output_map` was present. Update `PjRtCompatibleClientRemapArrays` to construct output buffer assignments directly from `plan.input_devices_for_output_map()` when available, bypassing `plan.mappings()` traversal. Enforce in `RemapPlan::Validate` that `input_devices_for_output_map`, if specified, must be fully populated (`size() == num_outputs`). PiperOrigin-RevId: 974584311 --- third_party/xla/xla/python/ifrt/remap_plan.cc | 8 + .../xla/xla/python/ifrt/remap_plan_test.cc | 23 +++ .../xla/xla/python/pjrt_ifrt/pjrt_remap.cc | 184 +++++++++++++++--- 3 files changed, 187 insertions(+), 28 deletions(-) diff --git a/third_party/xla/xla/python/ifrt/remap_plan.cc b/third_party/xla/xla/python/ifrt/remap_plan.cc index e0e43c814425e4..4ea7e34451f3ce 100644 --- a/third_party/xla/xla/python/ifrt/remap_plan.cc +++ b/third_party/xla/xla/python/ifrt/remap_plan.cc @@ -576,6 +576,14 @@ absl::Status RemapPlan::Validate() const { } } + if (!rep_->input_devices_for_output_map.empty() && + rep_->input_devices_for_output_map.size() != num_outputs) { + return InvalidArgument( + "`input_devices_for_output_map` has %d outputs, but expected %d " + "outputs", + rep_->input_devices_for_output_map.size(), num_outputs); + } + for (int i = 0; i < num_outputs; ++i) { xla::ifrt::DeviceList* devices = rep_->output_specs[i].sharding->devices()->AddressableDeviceList(); diff --git a/third_party/xla/xla/python/ifrt/remap_plan_test.cc b/third_party/xla/xla/python/ifrt/remap_plan_test.cc index 83bfe931cb90e5..a32fe7ffd867c0 100644 --- a/third_party/xla/xla/python/ifrt/remap_plan_test.cc +++ b/third_party/xla/xla/python/ifrt/remap_plan_test.cc @@ -735,6 +735,29 @@ TEST_P(RemapPlanTest, InvalidInputDevicesForOutputMap) { HasSubstr("does not reference that device"))); } + { + std::vector multi_output_specs = {dummy_spec, dummy_spec}; + std::vector multi_mappings = { + RemapPlan::Mapping{/*in_array=*/0, + /*out_array=*/0, + /*from=*/{RemapPlan::Interval{0, 2, 1}}, + /*to=*/{RemapPlan::Interval{0, 2, 1}}}, + RemapPlan::Mapping{/*in_array=*/1, + /*out_array=*/1, + /*from=*/{RemapPlan::Interval{0, 2, 1}}, + /*to=*/{RemapPlan::Interval{0, 2, 1}}}}; + absl::flat_hash_map> + input_devices_for_output_map; + input_devices_for_output_map.insert( + {0, {{0, dummy_spec.sharding->devices()}}}); + RemapPlan plan(input_specs, multi_output_specs, multi_mappings, + std::move(input_devices_for_output_map)); + EXPECT_THAT(plan.Validate(), + absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("has 1 outputs, but expected 2 outputs"))); + } + ASSERT_OK(RemapPlan::CreateOptimized(client(), std::move(input_specs), std::move(output_specs), std::move(mappings)) diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_remap.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_remap.cc index 9ae4d055a26570..bbddb6ed5934c8 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_remap.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_remap.cc @@ -17,9 +17,11 @@ limitations under the License. #include #include +#include #include #include +#include "absl/container/flat_hash_map.h" #include "absl/log/check.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" @@ -27,15 +29,15 @@ limitations under the License. #include "xla/pjrt/pjrt_client.h" #include "xla/pjrt/pjrt_layout.h" #include "xla/python/ifrt/array.h" +#include "xla/python/ifrt/device.h" +#include "xla/python/ifrt/device_list.h" #include "xla/python/ifrt/dtype.h" #include "xla/python/ifrt/remap_plan.h" #include "xla/python/ifrt/rtti.h" #include "xla/python/ifrt/shape.h" #include "xla/python/pjrt_ifrt/pjrt_array.h" #include "xla/tsl/concurrency/ref_count.h" -#include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/logging.h" -#include "xla/tsl/platform/statusor.h" #include "xla/util.h" namespace xla { @@ -97,34 +99,160 @@ PjRtCompatibleClientRemapArrays(PjRtCompatibleClient* client, ->size()); } - for (const RemapPlan::Mapping& mapping : plan.mappings()) { - ABSL_ASSIGN_OR_RETURN( - absl::Span> in_buffers, - static_cast(arrays[mapping.in_array].get()) - ->mutable_pjrt_buffers()); - PjRtArray::PjRtBuffers& out_buffers = out_buffers_list[mapping.out_array]; - for (int s = 0; s < mapping.from.size(); ++s) { - const RemapPlan::Interval& in_interval = mapping.from[s]; - const RemapPlan::Interval& out_interval = mapping.to[s]; - int64_t in_shard = in_interval.start; - int64_t out_shard = out_interval.start; - switch (semantics) { - case ArrayCopySemantics::kReuseInput: - while (in_shard < in_interval.end) { - out_buffers[out_shard] = in_buffers[in_shard]; - in_shard += in_interval.step; - out_shard += out_interval.step; + // Handle outputs using `input_devices_for_output_map` when specified. + if (!plan.input_devices_for_output_map().empty()) { + std::vector>> + in_device_to_shard(num_inputs); + + for (const auto& [out_array, input_ranges] : + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) + plan.input_devices_for_output_map()) { + if (out_array < 0 || out_array >= num_outputs) { + return InvalidArgument("out_array must be in [0, %d), but is %d", + num_outputs, out_array); + } + absl::Span out_devices = plan.output_specs()[out_array] + .sharding->devices() + ->AddressableDeviceList() + ->devices(); + // Skip outputs with no addressable shards on this controller. + if (out_devices.empty()) { + continue; + } + + absl::flat_hash_map device_to_out_shard; + device_to_out_shard.reserve(out_devices.size()); + for (int out_shard = 0; out_shard < out_devices.size(); ++out_shard) { + device_to_out_shard.insert({out_devices[out_shard], out_shard}); + } + + PjRtArray::PjRtBuffers& out_buffers = out_buffers_list[out_array]; + + for (const RemapPlan::InputDeviceRange& input_range : input_ranges) { + absl::Span range_devices = + input_range.input_devices->AddressableDeviceList()->devices(); + // Skip input ranges with no addressable shards on this controller. + if (range_devices.empty()) { + continue; + } + + int in_array = input_range.in_array; + if (in_array < 0 || in_array >= num_inputs) { + return InvalidArgument("in_array must be in [0, %d), but is %d", + num_inputs, in_array); + } + ABSL_ASSIGN_OR_RETURN( + absl::Span> in_buffers, + static_cast(arrays[in_array].get()) + ->mutable_pjrt_buffers()); + + if (!in_device_to_shard[in_array].has_value()) { + absl::Span in_devices = arrays[in_array] + ->sharding() + .devices() + ->AddressableDeviceList() + ->devices(); + auto& map = in_device_to_shard[in_array].emplace(); + map.reserve(in_devices.size()); + for (int in_shard = 0; in_shard < in_devices.size(); ++in_shard) { + map.insert({in_devices[in_shard], in_shard}); + } + } + + for (Device* const device : range_devices) { + auto in_it = in_device_to_shard[in_array]->find(device); + if (in_it == in_device_to_shard[in_array]->end()) { + return InvalidArgument("Device %s not found in input array %d", + device->DebugString(), in_array); } - break; - case ArrayCopySemantics::kDonateInput: - while (in_shard < in_interval.end) { - out_buffers[out_shard] = std::move(in_buffers[in_shard]); - in_shard += in_interval.step; - out_shard += out_interval.step; + int in_shard = in_it->second; + + auto out_it = device_to_out_shard.find(device); + if (out_it == device_to_out_shard.end()) { + return InvalidArgument("Device %s not found in output array %d", + device->DebugString(), out_array); + } + int out_shard = out_it->second; + + if (out_buffers[out_shard] != nullptr) { + return InvalidArgument( + "Output array %d shard on device %s is assigned more than once", + out_array, device->DebugString()); } - break; - default: - return InvalidArgument("Invalid ArrayCopySemantics: %d", semantics); + + switch (semantics) { + case ArrayCopySemantics::kReuseInput: + out_buffers[out_shard] = in_buffers[in_shard]; + break; + case ArrayCopySemantics::kDonateInput: + if (in_buffers[in_shard] == nullptr) { + return InvalidArgument( + "Input array %d shard on device %s is used more than once", + in_array, device->DebugString()); + } + out_buffers[out_shard] = std::move(in_buffers[in_shard]); + break; + default: + return InvalidArgument("Invalid ArrayCopySemantics: %d", + semantics); + } + } + } + } + } else { + for (const RemapPlan::Mapping& mapping : plan.mappings()) { + ABSL_ASSIGN_OR_RETURN( + absl::Span> in_buffers, + static_cast(arrays[mapping.in_array].get()) + ->mutable_pjrt_buffers()); + PjRtArray::PjRtBuffers& out_buffers = out_buffers_list[mapping.out_array]; + for (int s = 0; s < mapping.from.size(); ++s) { + const RemapPlan::Interval& in_interval = mapping.from[s]; + const RemapPlan::Interval& out_interval = mapping.to[s]; + int64_t in_shard = in_interval.start; + int64_t out_shard = out_interval.start; + switch (semantics) { + case ArrayCopySemantics::kReuseInput: + while (in_shard < in_interval.end) { + if (out_buffers[out_shard] != nullptr) { + return InvalidArgument( + "Output array %d shard %d is assigned more than once", + mapping.out_array, out_shard); + } + out_buffers[out_shard] = in_buffers[in_shard]; + in_shard += in_interval.step; + out_shard += out_interval.step; + } + break; + case ArrayCopySemantics::kDonateInput: + while (in_shard < in_interval.end) { + if (out_buffers[out_shard] != nullptr) { + return InvalidArgument( + "Output array %d shard %d is assigned more than once", + mapping.out_array, out_shard); + } + if (in_buffers[in_shard] == nullptr) { + return InvalidArgument( + "Input array %d shard %d is used more than once", + mapping.in_array, in_shard); + } + out_buffers[out_shard] = std::move(in_buffers[in_shard]); + in_shard += in_interval.step; + out_shard += out_interval.step; + } + break; + default: + return InvalidArgument("Invalid ArrayCopySemantics: %d", semantics); + } + } + } + } + + for (int i = 0; i < num_outputs; ++i) { + for (int s = 0; s < out_buffers_list[i].size(); ++s) { + if (out_buffers_list[i][s] == nullptr) { + return InvalidArgument( + "Output array %d addressable shard %d is unassigned", i, s); } } } From 7bd38880cece6eb7f9e94fc2a0357bae31b26563 Mon Sep 17 00:00:00 2001 From: Tori Baker Date: Tue, 1 Sep 2026 10:49:41 -0700 Subject: [PATCH 16/16] Replace TF_ASSERT_OK_AND_ASSIGN with ASSERT_OK_AND_ASSIGN in triton_fusion_analysis_test.cc PiperOrigin-RevId: 974585794 --- .../gpu/triton_fusion_analysis_test.cc | 146 +++++++++--------- 1 file changed, 73 insertions(+), 73 deletions(-) diff --git a/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc b/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc index cde00093180611..74325b23ede671 100644 --- a/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc +++ b/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc @@ -42,8 +42,8 @@ using ::testing::FieldsAre; using TritonDotAnalysisTest = HloHardwareIndependentTestBase; TEST_F(TritonDotAnalysisTest, QueryingOutputScopeParametersAlwaysWorks) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(R"( + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( triton_dot { p0 = f32[8,8] parameter(0) ROOT dot = f32[8,8] dot(p0, p0), @@ -54,7 +54,7 @@ ENTRY e { p0 = f32[8,8] parameter(0) ROOT r = f32[8,8] fusion(p0), kind=kCustom, calls=triton_dot })")); - TF_ASSERT_OK_AND_ASSIGN( + ASSERT_OK_AND_ASSIGN( const auto analysis, TritonFusionAnalysis::Execute(*module->entry_computation() ->root_instruction() @@ -85,16 +85,16 @@ ENTRY e { called_computations={triton_dot} ROOT bitcast.2 = bf16[1,8,6,3]{3,2,1,0} bitcast(custom-call) })"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_text)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); const HloComputation* dot_computation = module->entry_computation() ->root_instruction() ->operand(0) ->called_computations()[0]; const HloInstruction* p0 = dot_computation->parameter_instruction(0); const HloInstruction* p1 = dot_computation->parameter_instruction(1); - TF_ASSERT_OK_AND_ASSIGN(const auto analysis, - TritonFusionAnalysis::Execute(*dot_computation)); + ASSERT_OK_AND_ASSIGN(const auto analysis, + TritonFusionAnalysis::Execute(*dot_computation)); EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(), p0); EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::RHS).begin(), @@ -142,14 +142,14 @@ ENTRY e { "num_stages":1,"num_warps":2, "num_ctas":1}}} })"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_text)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); const HloComputation* dot_computation = module->entry_computation()->root_instruction()->called_computations()[0]; const HloInstruction* p0 = dot_computation->parameter_instruction(0); const HloInstruction* p1 = dot_computation->parameter_instruction(1); - TF_ASSERT_OK_AND_ASSIGN(const auto analysis, - TritonFusionAnalysis::Execute(*dot_computation)); + ASSERT_OK_AND_ASSIGN(const auto analysis, + TritonFusionAnalysis::Execute(*dot_computation)); EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(), p0); EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::RHS).begin(), @@ -197,16 +197,16 @@ ENTRY e { called_computations={triton_dot} ROOT bitcast.2 = bf16[1,8,6,3]{3,2,1,0} bitcast(custom-call) })"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_text)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); const HloComputation* dot_computation = module->entry_computation() ->root_instruction() ->operand(0) ->called_computations()[0]; const HloInstruction* p0 = dot_computation->parameter_instruction(0); const HloInstruction* p1 = dot_computation->parameter_instruction(1); - TF_ASSERT_OK_AND_ASSIGN(const auto analysis, - TritonFusionAnalysis::Execute(*dot_computation)); + ASSERT_OK_AND_ASSIGN(const auto analysis, + TritonFusionAnalysis::Execute(*dot_computation)); EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(), p0); EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::RHS).begin(), @@ -253,14 +253,14 @@ ENTRY e { custom_call_target="__triton", called_computations={triton_dot} })"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_text)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); const HloComputation* dot_computation = module->entry_computation()->root_instruction()->called_computations()[0]; const HloInstruction* p0 = dot_computation->parameter_instruction(0); const HloInstruction* p1 = dot_computation->parameter_instruction(1); - TF_ASSERT_OK_AND_ASSIGN(const auto analysis, - TritonFusionAnalysis::Execute(*dot_computation)); + ASSERT_OK_AND_ASSIGN(const auto analysis, + TritonFusionAnalysis::Execute(*dot_computation)); EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(), p1); EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::RHS).begin(), @@ -309,16 +309,16 @@ ENTRY e { called_computations={triton_dot} ROOT bitcast.2 = bf16[1,8,6,3]{3,2,1,0} bitcast(custom-call) })"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_text)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); const HloComputation* dot_computation = module->entry_computation() ->root_instruction() ->operand(0) ->called_computations()[0]; const HloInstruction* p0 = dot_computation->parameter_instruction(0); const HloInstruction* p1 = dot_computation->parameter_instruction(1); - TF_ASSERT_OK_AND_ASSIGN(const auto analysis, - TritonFusionAnalysis::Execute(*dot_computation)); + ASSERT_OK_AND_ASSIGN(const auto analysis, + TritonFusionAnalysis::Execute(*dot_computation)); EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(), p0); EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::RHS).begin(), @@ -368,16 +368,16 @@ ENTRY e { called_computations={triton_dot} ROOT bitcast.2 = bf16[1,8,6,3]{3,2,1,0} bitcast(custom-call) })"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_text)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); const HloComputation* dot_computation = module->entry_computation() ->root_instruction() ->operand(0) ->called_computations()[0]; const HloInstruction* p0 = dot_computation->parameter_instruction(0); const HloInstruction* p1 = dot_computation->parameter_instruction(1); - TF_ASSERT_OK_AND_ASSIGN(const auto analysis, - TritonFusionAnalysis::Execute(*dot_computation)); + ASSERT_OK_AND_ASSIGN(const auto analysis, + TritonFusionAnalysis::Execute(*dot_computation)); EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(), p0); EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::RHS).begin(), @@ -424,16 +424,16 @@ ENTRY e { custom_call_target="__triton", called_computations={triton_dot} ROOT bitcast.2 = bf16[3,8,1,3]{3,2,1,0} bitcast(custom-call) })"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_text)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); const HloComputation* dot_computation = module->entry_computation() ->root_instruction() ->operand(0) ->called_computations()[0]; const HloInstruction* p0 = dot_computation->parameter_instruction(0); const HloInstruction* p1 = dot_computation->parameter_instruction(1); - TF_ASSERT_OK_AND_ASSIGN(const auto analysis, - TritonFusionAnalysis::Execute(*dot_computation)); + ASSERT_OK_AND_ASSIGN(const auto analysis, + TritonFusionAnalysis::Execute(*dot_computation)); EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(), p0); EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::RHS).begin(), @@ -483,13 +483,13 @@ ENTRY e { ROOT r = bf16[3,12,2]{2,1,0} fusion(p0, p1), kind=kCustom, calls=triton_dot })"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_text)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); const HloComputation* dot_computation = module->entry_computation()->root_instruction()->called_computations()[0]; const HloInstruction* dot_output = dot_computation->root_instruction(); - TF_ASSERT_OK_AND_ASSIGN(const auto analysis, - TritonFusionAnalysis::Execute(*dot_computation)); + ASSERT_OK_AND_ASSIGN(const auto analysis, + TritonFusionAnalysis::Execute(*dot_computation)); EXPECT_THAT( *analysis.IterSpec(TritonFusionAnalysis::Scope::OUTPUT, dot_output, 0), ElementsAre(FieldsAre(/*stride=*/1, /*count=*/24, /*slice_start=*/0, @@ -505,8 +505,8 @@ ENTRY e { } TEST_F(TritonDotAnalysisTest, OutputParameterIsHandled) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(R"( + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( HloModule t triton_dot { @@ -531,8 +531,8 @@ ENTRY e { module->entry_computation()->root_instruction()->called_computations()[0]; const HloInstruction* output_param = dot_computation->parameter_instruction(2); - TF_ASSERT_OK_AND_ASSIGN(const auto analysis, - TritonFusionAnalysis::Execute(*dot_computation)); + ASSERT_OK_AND_ASSIGN(const auto analysis, + TritonFusionAnalysis::Execute(*dot_computation)); EXPECT_EQ( analysis.IterSpec(TritonFusionAnalysis::Scope::OUTPUT, output_param, 0) ->size(), @@ -556,8 +556,8 @@ ENTRY e { } TEST_F(TritonDotAnalysisTest, InputBroadcastFromScalarIsHandled) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(R"( + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( HloModule t triton_dot { @@ -577,8 +577,8 @@ ENTRY e { const HloComputation* dot_computation = module->entry_computation()->root_instruction()->called_computations()[0]; const HloInstruction* scalar = dot_computation->parameter_instruction(1); - TF_ASSERT_OK_AND_ASSIGN(const auto analysis, - TritonFusionAnalysis::Execute(*dot_computation)); + ASSERT_OK_AND_ASSIGN(const auto analysis, + TritonFusionAnalysis::Execute(*dot_computation)); EXPECT_EQ(analysis.IterSpec(TritonFusionAnalysis::Scope::RHS, scalar, 0), nullptr); EXPECT_EQ(analysis.IterSpec(TritonFusionAnalysis::Scope::RHS, scalar, 1), @@ -586,8 +586,8 @@ ENTRY e { } TEST_F(TritonDotAnalysisTest, InputBroadcastFromVectorIsHandled) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(R"( + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( HloModule t triton_dot { @@ -607,8 +607,8 @@ ENTRY e { const HloComputation* dot_computation = module->entry_computation()->root_instruction()->called_computations()[0]; const HloInstruction* vector = dot_computation->parameter_instruction(1); - TF_ASSERT_OK_AND_ASSIGN(const auto analysis, - TritonFusionAnalysis::Execute(*dot_computation)); + ASSERT_OK_AND_ASSIGN(const auto analysis, + TritonFusionAnalysis::Execute(*dot_computation)); EXPECT_EQ( analysis.IterSpec(TritonFusionAnalysis::Scope::RHS, vector, 0)->size(), 1); @@ -620,8 +620,8 @@ ENTRY e { } TEST_F(TritonDotAnalysisTest, BroadcastFromTriviallySizedDimensionIsSupported) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(R"( + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( f { p0 = f16[2] parameter(0) bc0 = f16[1,2] bitcast(p0) @@ -641,8 +641,8 @@ e { ->root_instruction() ->called_computations()[0]; const HloInstruction* p0 = dot_computation.parameter_instruction(0); - TF_ASSERT_OK_AND_ASSIGN(const auto analysis, - TritonFusionAnalysis::Execute(dot_computation)); + ASSERT_OK_AND_ASSIGN(const auto analysis, + TritonFusionAnalysis::Execute(dot_computation)); EXPECT_EQ(analysis.IterSpec(TritonFusionAnalysis::Scope::LHS, p0, 0)->size(), 1); EXPECT_THAT(*analysis.IterSpec(TritonFusionAnalysis::Scope::LHS, p0, 0), @@ -658,8 +658,8 @@ e { } TEST_F(TritonDotAnalysisTest, BroadcastWithinDimensionIsNotSupported) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(R"( + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( f { a = f16[5,7] parameter(0) br = f16[2,5,7] broadcast(a), dimensions={1,2} @@ -684,8 +684,8 @@ e { } TEST_F(TritonDotAnalysisTest, OutputBroadcastIsNotAccepted) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(R"( + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( HloModule t ENTRY e { @@ -705,8 +705,8 @@ ENTRY e { } TEST_F(TritonDotAnalysisTest, DegenerateSplitFragmentIsHandled) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(R"( + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( triton_gemm_r { Arg_0.1 = s8[30,913,8,21]{3,2,1,0} parameter(0) bitcast.6 = s8[30,8,21,913]{2,1,3,0} bitcast(Arg_0.1) @@ -730,8 +730,8 @@ ENTRY e { })")); const HloComputation* dot_computation = module->entry_computation()->root_instruction()->called_computations()[0]; - TF_ASSERT_OK_AND_ASSIGN(const auto analysis, - TritonFusionAnalysis::Execute(*dot_computation)); + ASSERT_OK_AND_ASSIGN(const auto analysis, + TritonFusionAnalysis::Execute(*dot_computation)); EXPECT_THAT(*analysis.IterSpec(TritonFusionAnalysis::Scope::OUTPUT, dot_computation->root_instruction(), 0), ElementsAre(FieldsAre(/*stride=*/1, /*count=*/8 * 21, @@ -748,8 +748,8 @@ TEST_F(TritonDotAnalysisTest, HandlesFurtherPropagationFromTrivialSizedTensorGracefully) { // We could probably support this better, just checking to avoid a crash for // now. - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(R"( + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( triton_gemm_r { a = f32[3,3]{1,0} parameter(0) constant = f32[1,1]{1,0} constant({ {0} }) @@ -777,8 +777,8 @@ ENTRY e { } TEST_F(TritonDotAnalysisTest, DynamicSliceIsSupported) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(R"( + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( triton_gemm { dot_lhs = f32[2,18]{1,0} parameter(0) dynamic_slice_input = f32[96,2]{1,0} parameter(1) @@ -806,8 +806,8 @@ ENTRY e { const HloComputation* dot_computation = module->entry_computation()->root_instruction()->called_computations()[0]; - TF_ASSERT_OK_AND_ASSIGN(const auto analysis, - TritonFusionAnalysis::Execute(*dot_computation)); + ASSERT_OK_AND_ASSIGN(const auto analysis, + TritonFusionAnalysis::Execute(*dot_computation)); const HloInstruction* p0 = dot_computation->parameter_instruction(0); const HloInstruction* p1 = dot_computation->parameter_instruction(1); EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(), @@ -837,8 +837,8 @@ ENTRY e { } TEST_F(TritonDotAnalysisTest, QueryScopeAlwaysWorks) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(R"( + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( triton_gemm_r { Arg_0.1 = s8[30,913,8,21]{3,2,1,0} parameter(0) bitcast.6 = s8[30,8,21,913]{2,1,3,0} bitcast(Arg_0.1) @@ -862,8 +862,8 @@ ENTRY e { })")); const HloComputation* dot_computation = module->entry_computation()->root_instruction()->called_computations()[0]; - TF_ASSERT_OK_AND_ASSIGN(const auto analysis, - TritonFusionAnalysis::Execute(*dot_computation)); + ASSERT_OK_AND_ASSIGN(const auto analysis, + TritonFusionAnalysis::Execute(*dot_computation)); for (const auto& hlo : dot_computation->instructions()) { if (hlo->opcode() != HloOpcode::kDot) { EXPECT_TRUE(analysis.QueryInstructionScope(*hlo).has_value()); @@ -872,8 +872,8 @@ ENTRY e { } TEST_F(TritonDotAnalysisTest, ScaledDotIsSupported) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(R"( + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( HloModule test scaled_dot { @@ -886,8 +886,8 @@ scaled_dot { rhs_batch_dims={0}, rhs_contracting_dims={1} })")); const HloComputation* dot_computation = *module->computations().begin(); - TF_ASSERT_OK_AND_ASSIGN(const auto analysis, - TritonFusionAnalysis::Execute(*dot_computation)); + ASSERT_OK_AND_ASSIGN(const auto analysis, + TritonFusionAnalysis::Execute(*dot_computation)); const HloInstruction* lhs = dot_computation->parameter_instruction(0); const HloInstruction* rhs = dot_computation->parameter_instruction(1); const HloInstruction* lhs_scale = dot_computation->parameter_instruction(2);