Merge pull request #1025 from dlee992/skip_if_meet_raise_bb

skip `raise` basic blocks in `verifyFanoutBackward`
numba · Apr 17, 2024 · 622c33f · 622c33f
2 parents 42b9834 + 01d6aed
commit 622c33f
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 4 deletions.
diff --git a/ffi/custom_passes.cpp b/ffi/custom_passes.cpp
@@ -400,7 +400,8 @@ struct RefPrunePass : public FunctionPass {
                     SmallBBSet tail_nodes;
                     tail_nodes.insert(decref->getParent());
                     if (!verifyFanoutBackward(incref, incref->getParent(),
-                                              &tail_nodes))
+                                              &tail_nodes, false))
+
                         continue;
 
                     // scan the CFG between the incref and decref BBs, if
@@ -495,6 +496,29 @@ struct RefPrunePass : public FunctionPass {
      *   │ MORE CFG   │
      *   └────────────┘
      *
+     * a complex pattern about fanout-raise
+     * https://github.com/numba/llvmlite/issues/1023
+     *           ┌────────────┐
+     *           │   incref   │
+     *           │   incref   │
+     *           └────────────┘
+     *             /           \
+     *            /             \
+     *     ┌────────────┐        \
+     *     │   decref   |         \
+     *     └────────────┘          \
+     *      /          \            \
+     *     /            \            \
+     * ┌────────────┐ ┌────────────┐  \
+     * │   decref   | │   incref   |   \
+     * └────────────┘ └────────────┘    \
+     *                 /            \    \
+     *                /              \    \
+     *           ┌────────────┐      ┌────────────┐
+     *           │   decref   |      │   raise    |
+     *           │   decref   |      └────────────┘
+     *           └────────────┘
+     *
      * Parameters:
      *  - F a Function
      *  - prune_raise_exit, if false case 1 is considered, if true case 2 is
@@ -648,10 +672,12 @@ struct RefPrunePass : public FunctionPass {
                 for (BasicBlock *bb : *decref_blocks) {
                     raising_blocks.insert(bb);
                 }
-                if (verifyFanoutBackward(incref, head_node, p_raising_blocks))
+                if (verifyFanoutBackward(incref, head_node, p_raising_blocks,
+                                         prune_raise_exit))
                     return true;
 
-            } else if (verifyFanoutBackward(incref, head_node, decref_blocks)) {
+            } else if (verifyFanoutBackward(incref, head_node, decref_blocks,
+                                            prune_raise_exit)) {
                 return true;
             }
         }
@@ -844,7 +870,8 @@ struct RefPrunePass : public FunctionPass {
      *
      */
     bool verifyFanoutBackward(CallInst *incref, BasicBlock *head_node,
-                              const SmallBBSet *tail_nodes) {
+                              const SmallBBSet *tail_nodes,
+                              bool prune_raise_exit) {
         // push the tail nodes into a work list
         SmallVector<BasicBlock *, 10> todo;
         for (BasicBlock *bb : *tail_nodes) {
@@ -864,6 +891,10 @@ struct RefPrunePass : public FunctionPass {
             while (workstack.size() > 0) {
                 // Get a basic block
                 BasicBlock *cur_node = workstack.pop_back_val();
+                // If cur_node is a raising block, then skip it
+                if (prune_raise_exit && isRaising(cur_node)) {
+                    continue;
+                }
                 // if the block has been seen before then skip
                 if (visited.count(cur_node)) {
                     // Already visited

diff --git a/llvmlite/tests/test_refprune.py b/llvmlite/tests/test_refprune.py
@@ -521,6 +521,37 @@ def test_fanout_raise_5(self):
         mod, stats = self.check(self.fanout_raise_5)
         self.assertEqual(stats.fanout_raise, 2)
 
+    # test case 6 is from https://github.com/numba/llvmlite/issues/1023
+    fanout_raise_6 = r"""
+define i32 @main(i8* %ptr, i1 %cond1, i1 %cond2, i1 %cond3, i8** %excinfo) {
+bb_A:
+    call void @NRT_incref(i8* %ptr)
+    call void @NRT_incref(i8* %ptr)
+    br i1 %cond1, label %bb_B, label %bb_C
+bb_B:
+    call void @NRT_decref(i8* %ptr)
+    br i1 %cond2, label %bb_D, label %bb_E
+bb_C:
+    store i8* null, i8** %excinfo, !numba_exception_output !0
+    ret i32 1
+bb_D:
+    call void @NRT_decref(i8* %ptr)
+    ret i32 0
+bb_E:
+    call void @NRT_incref(i8* %ptr)
+    br i1 %cond3, label %bb_F, label %bb_C
+bb_F:
+    call void @NRT_decref(i8* %ptr)
+    call void @NRT_decref(i8* %ptr)
+    ret i32 0
+}
+!0 = !{i1 1}
+"""
+
+    def test_fanout_raise_6(self):
+        mod, stats = self.check(self.fanout_raise_6)
+        self.assertEqual(stats.fanout_raise, 7)
+
 
 if __name__ == '__main__':
     unittest.main()