[Mlir-commits] [mlir] [mlir][dataflow] Remove early exit in dead code analysis for zero-operand returns (PR #68151)

[Jacob Mai Peng]

https://github.com/pengmai created https://github.com/llvm/llvm-project/pull/68151

The PredecessorState in dead code analysis does not register zero-operand returns as predecessors of their corresponding call ops. This causes bugs with dense dataflow analyses that use dead code analysis to query for the predecessors of CallOpInterfaces.

>From 2f0dadc0da24fa7303d0fa7103e9ea5100501bf3 Mon Sep 17 00:00:00 2001
From: Jacob Peng <jacobmpeng at gmail.com>
Date: Tue, 3 Oct 2023 15:55:12 -0400
Subject: [PATCH] [mlir][dataflow] Remove early exit in dead code analysis for
 zero-operand returns

The PredecessorState in dead code analysis does not register zero-operand returns as predecessors of their corresponding call ops. This causes bugs with dense analyses that use dead code analysis to query for the predecessors of CallOpInterfaces.
---
 mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp |  4 ----
 .../DataFlow/test-dead-code-analysis.mlir       | 12 ++++++++++++
 .../DataFlow/test-last-modified-callgraph.mlir  | 17 +++++++++++++++++
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp b/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
index d423d37b9770c69..bbe139fd4fb69fe 100644
--- a/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/DeadCodeAnalysis.cpp
@@ -407,10 +407,6 @@ void DeadCodeAnalysis::visitRegionTerminator(
 
 void DeadCodeAnalysis::visitCallableTerminator(Operation *op,
                                                CallableOpInterface callable) {
-  // If there are no exiting values, we have nothing to do.
-  if (op->getNumOperands() == 0)
-    return;
-
   // Add as predecessors to all callsites this return op.
   auto *callsites = getOrCreateFor<PredecessorState>(op, callable);
   bool canResolve = op->hasTrait<OpTrait::ReturnLike>();
diff --git a/mlir/test/Analysis/DataFlow/test-dead-code-analysis.mlir b/mlir/test/Analysis/DataFlow/test-dead-code-analysis.mlir
index 020c8b2a0e84dc7..b1af670440c2d97 100644
--- a/mlir/test/Analysis/DataFlow/test-dead-code-analysis.mlir
+++ b/mlir/test/Analysis/DataFlow/test-dead-code-analysis.mlir
@@ -118,6 +118,18 @@ func.func @test_callgraph(%cond: i1, %arg0: i32) -> i32 {
   return %2 : i32
 }
 
+func.func private @bax(%arg0: i32) {
+  return {void_return}
+}
+
+func.func @test_callgraph_void_return(%arg0: i32) -> i32 {
+  // CHECK: call_void_return:
+  // CHECK: op_preds: (all) predecessors:
+  // CHECK:   func.return {void_return}
+  func.call @bax(%arg0) {tag = "call_void_return"}: (i32) -> ()
+  return %arg0 : i32
+}
+
 // CHECK: test_unknown_branch:
 // CHECK:  region #0
 // CHECK:   ^bb0 = live
diff --git a/mlir/test/Analysis/DataFlow/test-last-modified-callgraph.mlir b/mlir/test/Analysis/DataFlow/test-last-modified-callgraph.mlir
index 4a243571c231a19..709d787bb306bef 100644
--- a/mlir/test/Analysis/DataFlow/test-last-modified-callgraph.mlir
+++ b/mlir/test/Analysis/DataFlow/test-last-modified-callgraph.mlir
@@ -68,6 +68,23 @@ func.func @test_multiple_return_sites(%cond: i1, %a: i32, %ptr: memref<i32>) ->
 
 // -----
 
+// CHECK-LABEL: test_tag: after_call
+// CHECK: operand #0
+// CHECK-NEXT: - write0
+func.func private @void_return(%ptr: memref<i32>) {
+  return
+}
+
+func.func @test_call_void_return() {
+  %ptr = memref.alloc() : memref<i32>
+  %c0 = arith.constant 0 : i32
+  memref.store %c0, %ptr[] {tag_name = "write0"} : memref<i32>
+  func.call @void_return(%ptr) : (memref<i32>) -> ()
+  memref.load %ptr[] {tag = "after_call"} : memref<i32>
+  return
+}
+
+// -----
 
 func.func private @callee(%arg0: memref<f32>) -> memref<f32> {
   %2 = arith.constant 2.0 : f32