diff --git a/mlir/lib/Util/Dependency.cpp b/mlir/lib/Util/Dependency.cpp
index 4a301ceaf..677b0d05d 100644
--- a/mlir/lib/Util/Dependency.cpp
+++ b/mlir/lib/Util/Dependency.cpp
@@ -595,6 +595,8 @@ void addAsyncDependencyIfNew(Operation *op, Value token) {
     return;
   if (!token)
     return;
+  if (token.getDefiningOp() && token.getDefiningOp() == op)
+    return;
   if (auto async_op = dyn_cast<air::AsyncOpInterface>(op)) {
     addAsyncDependencyIfNewImpl(async_op, token);
   } else if (auto for_op = dyn_cast<scf::ForOp>(op)) {
@@ -674,8 +676,8 @@ bool isAsyncDependent(Operation *a, Operation *b) {
   for (auto dep : dep_b) {
     if (dep == token_a)
       return true;
-    else if (auto dep_wa_defop = dep.getDefiningOp<air::WaitAllOp>()) {
-      if (isAsyncDependent(a, dep_wa_defop))
+    else if (dep.getDefiningOp() && air::isAsyncOp(dep.getDefiningOp())) {
+      if (isAsyncDependent(a, dep.getDefiningOp()))
         return true;
     }
   }