From 8f23fac4da254e8cd2a3160a4fa029613a284ebe Mon Sep 17 00:00:00 2001
From: Eugene Zhulenev <ezhulenev@google.com>
Date: Thu, 27 May 2021 12:42:59 -0700
Subject: [PATCH] [mlir:Async] Convert assertions to async errors only inside
 async functions

Differential Revision: https://reviews.llvm.org/D103278
---
 mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp b/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
index 6ebf48ad6475..0789a0ee6887 100644
--- a/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
+++ b/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
@@ -591,8 +591,11 @@ void AsyncToAsyncRuntimePass::runOnOperation() {
   runtimeTarget.addIllegalOp<CreateGroupOp, AddToGroupOp>();
   runtimeTarget.addIllegalOp<ExecuteOp, AwaitOp, AwaitAllOp, async::YieldOp>();
 
-  // Assertions must be converted to runtime errors.
-  runtimeTarget.addIllegalOp<AssertOp>();
+  // Assertions must be converted to runtime errors inside async functions.
+  runtimeTarget.addDynamicallyLegalOp<AssertOp>([&](AssertOp op) -> bool {
+    auto func = op->getParentOfType<FuncOp>();
+    return outlinedFunctions.find(func) == outlinedFunctions.end();
+  });
   runtimeTarget.addLegalOp<CondBranchOp>();
 
   if (failed(applyPartialConversion(module, runtimeTarget,