[XLA] Verify that async computations are trivial (contain only a root and parameter instructions).

PiperOrigin-RevId: 566509445

Author: berkinilbeyi

xla/service/hlo_verifier.cc

@@ -2240,6 +2240,18 @@ Status VerifyAsynchronousInstructionPairs(const HloModule& module) {
   return OkStatus();
 }
 
+// Checks that the asynchronous computation only has a root and parameter
+// instructions.
+Status VerifyAsyncComputation(const HloComputation* async_computation) {
+  if (!async_computation->CanExpandIntoSingleInstruction()) {
+    return FailedPrecondition(
+        "Asynchronous computation %s expected to contain only the root and "
+        "parameter instructions.",
+        async_computation->name());
+  }
+  return OkStatus();
+}
+
 // Checks that AllReduce instructions in the module are either all layout
 // constrained or all unconstrained.
 Status VerifyLayoutConstrainedAllReduce(const HloModule& module) {
@@ -2849,6 +2861,9 @@ StatusOr<bool> HloVerifier::Run(
     for (auto* computation : module->computations(execution_threads)) {
       TF_RETURN_IF_ERROR(computation->Accept(shape_verifier.get()));
       TF_RETURN_IF_ERROR(computation->Accept(&instruction_verifier));
+      if (computation->IsAsyncComputation()) {
+        TF_RETURN_IF_ERROR(VerifyAsyncComputation(computation));
+      }
     }
 
     TF_RETURN_IF_ERROR(shape_verifier->VerifyEntryComputationLayout(*module));

xla/service/hlo_verifier_test.cc

@@ -1219,6 +1219,33 @@ TEST_F(HloVerifierTest, AsyncUpdateWrongType) {
           "async-update expects the shape of operand and output to match"));
 }
 
+TEST_F(HloVerifierTest, AsyncOpComputationNotTrivial) {
+  const char* const hlo_string = R"(
+  HloModule Module
+
+  async_computation {
+    p = f32[2,3] parameter(0)
+    copy = f32[2,3] copy(p)
+    ROOT custom-call = f32[3,2] custom-call(copy), custom_call_target="foo"
+  }
+
+  ENTRY AsyncStartAndAsyncDone {
+    p0 = f32[2,3] parameter(0)
+    async-start = ((f32[2,3]), f32[3,2], u32[]) async-start(p0), calls=async_computation
+    ROOT async-done = f32[3,2] async-done(async-start), calls=async_computation
+  }
+  )";
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          ParseAndReturnUnverifiedModule(hlo_string));
+
+  auto status = verifier().Run(module.get()).status();
+  ASSERT_FALSE(status.ok());
+  EXPECT_THAT(
+      status.message(),
+      HasSubstr(
+          "expected to contain only the root and parameter instructions"));
+}
+
 TEST_F(HloVerifierTestLayoutSensitive, AsyncDoneWrongGroupId) {
   const char* const hlo_string = R"(
   HloModule Module