From a81aac289df2c11b271d62b1b8a895ebc4684df0 Mon Sep 17 00:00:00 2001 From: Johannes Reifferscheid Date: Mon, 6 May 2024 01:19:13 -0700 Subject: [PATCH] Compute non-fusion computations just once. PiperOrigin-RevId: 630962946 --- third_party/xla/xla/service/gpu/priority_fusion.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/third_party/xla/xla/service/gpu/priority_fusion.cc b/third_party/xla/xla/service/gpu/priority_fusion.cc index 86db57c5212653..5fc09e0c36aeeb 100644 --- a/third_party/xla/xla/service/gpu/priority_fusion.cc +++ b/third_party/xla/xla/service/gpu/priority_fusion.cc @@ -697,8 +697,9 @@ absl::StatusOr GpuPriorityFusion::Run( // With this modification it will be easier to match instructions before and // after fusion passes, because they will have the same unique prefix. Names // are not used in the pipeline, but it makes debugging much easier. - for (auto* computation : - GetNonFusionComputations(module, execution_threads)) { + auto non_fusion_computations = + GetNonFusionComputations(module, execution_threads); + for (auto* computation : non_fusion_computations) { for (auto* instruction : computation->instructions()) { module->SetAndUniquifyInstrName(instruction, absl::StrCat(instruction->name(), ".0")); @@ -735,8 +736,7 @@ absl::StatusOr GpuPriorityFusion::Run( } int changed = false; - for (auto* computation : - GetNonFusionComputations(module, execution_threads)) { + for (auto* computation : non_fusion_computations) { if (computations_not_to_fuse.contains(computation)) continue; CHECK(!computation->IsFusionComputation());