PR #20313: Fix async wrapper to walk child computations

xla/service/gpu/transforms/BUILD

@@ -354,10 +354,12 @@ xla_cc_test(
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
         "//xla/hlo/testlib:verified_hlo_module",
+        "//xla/hlo/utils:hlo_query",
         "//xla/tests:hlo_test_base",
         "//xla/tests:literal_test_util",
         "@com_google_googletest//:gtest_main",
         "@tsl//tsl/platform:status_matchers",
+        "@tsl//tsl/platform:statusor",
         "@tsl//tsl/platform:test_main",
     ],
 )

xla/service/gpu/transforms/async_wrapper.cc

@@ -65,9 +65,10 @@ absl::StatusOr<bool> AsyncWrapper::Run(
         continue;
       }
 
-      // Otherwise, follow any `calls` to discover other instructions that can
-      // potentially be made async.
-      if (HloPredicateIsOp<HloOpcode::kCall>(instruction)) {
+      // Otherwise, follow anything other than `fusion`s to discover other
+      // instructions that can potentially be made async.
+      if (HloPredicateIsOp<HloOpcode::kCall, HloOpcode::kConditional,
+                           HloOpcode::kWhile>(instruction)) {
         std::copy(instruction->called_computations().begin(),
                   instruction->called_computations().end(),
                   std::back_inserter(computations));

xla/service/gpu/transforms/async_wrapper_test.cc

@@ -25,11 +25,13 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/pass/hlo_pass_interface.h"
 #include "xla/hlo/testlib/verified_hlo_module.h"
+#include "xla/hlo/utils/hlo_query.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
 #include "xla/tests/hlo_test_base.h"
 #include "xla/tests/literal_test_util.h"
 #include "tsl/platform/status_matchers.h"
+#include "tsl/platform/statusor.h"
 
 namespace xla::gpu {
 namespace {
@@ -81,5 +83,115 @@ TEST_F(AsyncWrapperTest, BasicFusion) {
   EXPECT_TRUE(LiteralTestUtil::Equal(expected, result));
 }
 
+TEST_F(AsyncWrapperTest, OpWithinWhileShouldWrapInAsync) {
+  const char* hlo = R"(
+  HloModule m
+
+  body {
+    param = (f32[1], s32[]) parameter(0)
+    p0 = f32[1] get-tuple-element(param), index=0
+    agg1 = f32[1] custom-call(p0), custom_call_target="foo"
+    agg2 = f32[1] custom-call(p0), custom_call_target="bar"
+    done = f32[1] add(agg1, agg2)
+    iter = s32[] get-tuple-element(param), index=1
+    c1 = s32[] constant(1)
+    add = s32[] add(iter, c1)
+    ROOT tuple = (f32[1], s32[]) tuple(done, add)
+  }
+
+  condition {
+    param.1 = (f32[1], s32[]) parameter(0)
+    iter.1 = s32[] get-tuple-element(param.1), index=1
+    c4 = s32[] constant(4)
+    ROOT compare = pred[] compare(iter.1, c4), direction=LT
+  }
+
+  ENTRY main {
+    c0 = s32[] constant(0)
+    p0.1 = f32[1] parameter(0)
+    agg3 = f32[1] custom-call(p0.1), custom_call_target="baz"
+    tuple = (f32[1], s32[]) tuple(agg3, c0)
+    while = (f32[1], s32[]) while(tuple), body=body, condition=condition
+    ROOT done.1 = f32[1] get-tuple-element(while), index=0
+  })";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          ParseAndReturnVerifiedModule(hlo));
+
+  AsyncWrapper wrapper(HloPredicateIsOp<HloOpcode::kCustomCall>);
+  TF_ASSERT_OK_AND_ASSIGN(bool changed,
+                          wrapper.Run(module.get(), /*execution_threads=*/{}));
+  EXPECT_TRUE(changed);
+  EXPECT_EQ(CountAsyncInstructions(module->entry_computation()), 2);
+  HloInstruction* while_op = hlo_query::FindInstruction(
+      module->entry_computation(), HloOpcode::kWhile);
+  ASSERT_NE(while_op, nullptr);
+  EXPECT_EQ(CountAsyncInstructions(while_op->while_body()), 4);
+}
+
+TEST_F(AsyncWrapperTest, OpWithinConditionalShouldWrapInAsync) {
+  const char* hlo = R"(
+  HloModule m
+
+  true_computation {
+    p0.1 = f32[] parameter(0)
+    ROOT res.1 = f32[] custom-call(p0.1), custom_call_target="foo"
+  }
+
+  false_computation {
+    p0.2 = f32[] parameter(0)
+    ROOT res.2 = f32[] custom-call(p0.2), custom_call_target="foo"
+  }
+
+  ENTRY main {
+    p0 = f32[] parameter(0)
+    c0 = f32[] constant(0)
+    compare = pred[] compare(p0, c0), direction=GE
+    ROOT done = f32[] conditional(compare, p0, p0), true_computation=true_computation, false_computation=false_computation
+  })";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          ParseAndReturnVerifiedModule(hlo));
+
+  AsyncWrapper wrapper(HloPredicateIsOp<HloOpcode::kCustomCall>);
+  TF_ASSERT_OK_AND_ASSIGN(bool changed,
+                          wrapper.Run(module.get(), /*execution_threads=*/{}));
+  EXPECT_TRUE(changed);
+  EXPECT_EQ(CountAsyncInstructions(module->entry_computation()), 0);
+  HloInstruction* conditional_op = hlo_query::FindInstruction(
+      module->entry_computation(), HloOpcode::kConditional);
+  ASSERT_NE(conditional_op, nullptr);
+  EXPECT_EQ(CountAsyncInstructions(conditional_op->true_computation()), 2);
+  EXPECT_EQ(CountAsyncInstructions(conditional_op->false_computation()), 2);
+}
+
+TEST_F(AsyncWrapperTest, OpWithinFusionShouldNotWrapInAsync) {
+  const char* hlo = R"(
+  foo {
+    p0 = f32[1] parameter(0)
+    ROOT custom-call = f32[1] custom-call(p0), custom_call_target="bar"
+  }
+  ENTRY main {
+    c0 = s32[] constant(0)
+    p0.1 = f32[1] parameter(0)
+    agg.1 = f32[1] fusion(p0.1), kind=kLoop, calls=foo
+    agg.2 = f32[1] custom-call(agg.1), custom_call_target="bar"
+    ROOT done.1 = (f32[1], f32[1]) tuple(agg.1, agg.2)
+  })";
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          ParseAndReturnVerifiedModule(hlo));
+
+  AsyncWrapper wrapper(HloPredicateIsOp<HloOpcode::kCustomCall>);
+  TF_ASSERT_OK_AND_ASSIGN(bool changed,
+                          wrapper.Run(module.get(), /*execution_threads=*/{}));
+  EXPECT_TRUE(changed);
+  EXPECT_EQ(CountAsyncInstructions(module->entry_computation()), 2);
+
+  HloInstruction* fusion = hlo_query::FindInstruction(
+      module->entry_computation(), HloOpcode::kFusion);
+  EXPECT_EQ(CountAsyncInstructions(fusion->fused_instructions_computation()),
+            0);
+}
+
 }  // namespace
 }  // namespace xla::gpu