[Mlir-commits] [mlir] [mlir][dataflow] Fix crash when call results mismatch predecessor return operands (PR #226086)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Divyansh Yadav (schizophrenicmaniac)

<details>
<summary>Changes</summary>

Fixes #<!-- -->226008

### Problem
When running `--sccp` (or sparse dataflow analysis) on call operations where the callee's return operation has a different number of operands than the call's forwarded results, `AbstractSparseForwardDataFlowAnalysis::visitCallOperation` asserts/crashes with:
```
Iteratees do not have equal size
```
This occurs, for example, when an `async.call` calls an `async.func`: the `async.func`'s `async.return` only provides operands for the explicitly returned values (e.g. `!async.value<i32>`), whereas the `async.call` produces both an implicit `!async.token` and the `!async.value<i32>`. The dataflow framework assumed a 1-to-1 correspondence between return operands and forwarded result lattices.

### Solution
- In `AbstractSparseForwardDataFlowAnalysis::visitCallOperation`, check if `predecessor->getNumOperands() != forwardedResultLattices.size()`. If they mismatch, conservatively set all forwarded result lattices to entry states and continue.
- Added corresponding defensive size checks across forward and backward sparse dataflow analyses (`visitCallableOperation` and backward call operand/result handling) and upgraded zips to `llvm::zip_equal`.
- Added a regression test in `mlir/test/Transforms/sccp-callgraph.mlir`.


---
Full diff: https://github.com/llvm/llvm-project/pull/226086.diff


2 Files Affected:

- (modified) mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp (+39-11) 
- (modified) mlir/test/Transforms/sccp-callgraph.mlir (+19) 


``````````diff
diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
index bc61c19832d29..029c7d6798639 100644
--- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
@@ -266,11 +266,21 @@ LogicalResult AbstractSparseForwardDataFlowAnalysis::visitCallOperation(
   ArrayRef<AbstractSparseLattice *> forwardedResultLattices =
       resultLattices.slice(firstForwarded, forwardedResults.size());
 
-  for (Operation *predecessor : predecessors->getKnownPredecessors())
+  for (Operation *predecessor : predecessors->getKnownPredecessors()) {
+    // If the predecessor (return-like op) has a different number of operands
+    // than the forwarded results, we can't directly relate them. This can
+    // happen when the callable implicitly produces some results (e.g.,
+    // async.func implicitly produces an async.token). Conservatively set
+    // all forwarded result lattices to entry states.
+    if (predecessor->getNumOperands() != forwardedResultLattices.size()) {
+      setAllToEntryStates(forwardedResultLattices);
+      continue;
+    }
     for (auto &&[operand, resLattice] :
          llvm::zip_equal(predecessor->getOperands(), forwardedResultLattices))
       join(resLattice,
            *getLatticeElementFor(getProgramPointAfter(call), operand));
+  }
   return success();
 }
 
@@ -288,10 +298,16 @@ void AbstractSparseForwardDataFlowAnalysis::visitCallableOperation(
   }
   for (Operation *callsite : callsites->getKnownPredecessors()) {
     auto call = cast<CallOpInterface>(callsite);
-    for (auto it : llvm::zip(call.getArgOperands(), argLattices))
-      join(std::get<1>(it),
-           *getLatticeElementFor(getProgramPointBefore(entryBlock),
-                                 std::get<0>(it)));
+    OperandRange argOperands = call.getArgOperands();
+    // If the number of forwarded call arguments doesn't match the callee's
+    // block arguments, we can't directly relate them. Conservatively set all
+    // argument lattices to entry states.
+    if (argOperands.size() != argLattices.size()) {
+      return setAllToEntryStates(argLattices);
+    }
+    for (auto [operand, argLattice] : llvm::zip_equal(argOperands, argLattices))
+      join(argLattice,
+           *getLatticeElementFor(getProgramPointBefore(entryBlock), operand));
   }
 }
 
@@ -542,11 +558,16 @@ AbstractSparseBackwardDataFlowAnalysis::visitOperation(Operation *op) {
       // Otherwise, propagate information from the entry point of the function
       // back to operands whenever possible.
       Block &block = region->front();
-      for (auto [blockArg, argOpOperand] :
-           llvm::zip(block.getArguments(), argOpOperands)) {
-        meet(getLatticeElement(argOpOperand.get()),
-             *getLatticeElementFor(getProgramPointAfter(op), blockArg));
-        unaccounted.reset(argOpOperand.getOperandNumber());
+      // If the number of forwarded operands doesn't match the callee's block
+      // arguments, we can't directly relate them. Fall through to treat all
+      // operands via visitCallOperand.
+      if (block.getNumArguments() == argOpOperands.size()) {
+        for (auto [blockArg, argOpOperand] :
+             llvm::zip_equal(block.getArguments(), argOpOperands)) {
+          meet(getLatticeElement(argOpOperand.get()),
+               *getLatticeElementFor(getProgramPointAfter(op), blockArg));
+          unaccounted.reset(argOpOperand.getOperandNumber());
+        }
       }
 
       // Handle the operands of the call op that aren't forwarded to any
@@ -605,7 +626,14 @@ LogicalResult AbstractSparseBackwardDataFlowAnalysis::visitCallableOperation(
           cast<CallOpInterface>(call).getForwardedResults();
       SmallVector<const AbstractSparseLattice *> callResultLattices =
           getLatticeElementsFor(getProgramPointAfter(op), forwardedResults);
-      for (auto [op, result] : llvm::zip(operandLattices, callResultLattices))
+      // If the number of return operands doesn't match the forwarded results,
+      // we can't directly relate them. Conservatively set all to exit states.
+      if (operandLattices.size() != callResultLattices.size()) {
+        setAllToExitStates(operandLattices);
+        continue;
+      }
+      for (auto [op, result] :
+           llvm::zip_equal(operandLattices, callResultLattices))
         meet(op, *result);
     }
   } else {
diff --git a/mlir/test/Transforms/sccp-callgraph.mlir b/mlir/test/Transforms/sccp-callgraph.mlir
index f31f749b6e75a..66e9914623e7b 100644
--- a/mlir/test/Transforms/sccp-callgraph.mlir
+++ b/mlir/test/Transforms/sccp-callgraph.mlir
@@ -306,3 +306,22 @@ func.func @bar(%arg0: index) -> index {
   // FUNC: cf.br ^bb1(%[[FOO]]
   cf.br ^bb1(%13 : index)
 }
+
+// -----
+
+/// Check that SCCP does not crash when callee returns fewer operands than the
+/// call's forwarded results (e.g., async.func implicitly produces a token).
+
+// CHECK-LABEL: async.func @async_func_implicit_token
+async.func @async_func_implicit_token() -> (!async.token, !async.value<i32>) {
+  %0 = arith.constant 42 : i32
+  return %0 : i32
+}
+
+// CHECK-LABEL: func @call_async_func_implicit_token
+func.func @call_async_func_implicit_token() {
+  // CHECK: async.call @async_func_implicit_token
+  %token, %val = async.call @async_func_implicit_token() : () -> (!async.token, !async.value<i32>)
+  return
+}
+

``````````

</details>


https://github.com/llvm/llvm-project/pull/226086