AllocOpt: Handle objref with no preserve_end in a single block. (#50277)

This fixes the AllocOpt pass when there is a single block and no
gc_preserve_end. The dominator tree traversal now considers the starting
(gc_preverse_begin) block as well and does not introduce a lifetime end
when there is no gc_preserve_end for the objref.

Author: Pangoraw

src/llvm-alloc-opt.cpp

@@ -433,36 +433,45 @@ void Optimizer::insertLifetime(Value *ptr, Constant *sz, Instruction *orig)
         abort();
     }
 #endif
-    // Record extra BBs that contain invisible uses.
+
+    // Record extra BBs that contain invisible uses with gc_preserve_{begin,end}.
+    // We traverse the dominator tree starting at each `gc_preserve_begin` and marking blocks
+    // as users until a corresponding `gc_preserve_end` is found. Blocks containing
+    // the `gc_preserve_end` have already been marked in the previous step.
     SmallSet<BasicBlock*, 8> extra_use;
     SmallVector<DomTreeNodeBase<BasicBlock>*, 8> dominated;
     for (auto preserve: use_info.preserves) {
-        for (auto RN = DT.getNode(preserve->getParent()); RN;
-             RN = dominated.empty() ? nullptr : dominated.pop_back_val()) {
-            for (auto N: *RN) {
-                auto bb = N->getBlock();
-                if (extra_use.count(bb))
-                    continue;
-                bool ended = false;
-                for (auto end: preserve->users()) {
-                    auto end_bb = cast<Instruction>(end)->getParent();
-                    auto end_node = DT.getNode(end_bb);
-                    if (end_bb == bb || (end_node && DT.dominates(end_node, N))) {
-                        ended = true;
-                        break;
-                    }
+        assert(dominated.empty());
+        dominated.push_back(DT.getNode(preserve->getParent()));
+        while (!dominated.empty()) {
+            auto N = dominated.pop_back_val();
+            if (!N) {
+                dominated.clear();
+                break;
+            }
+            auto bb = N->getBlock();
+            if (extra_use.count(bb))
+                continue;
+            bool ended = false;
+            for (auto end: preserve->users()) {
+                auto end_bb = cast<Instruction>(end)->getParent();
+                auto end_node = DT.getNode(end_bb);
+                if (end_bb == bb || (end_node && DT.dominates(end_node, N))) {
+                    ended = true;
+                    break;
                 }
-                if (ended)
-                    continue;
-                bbs.insert(bb);
-                extra_use.insert(bb);
-                dominated.push_back(N);
             }
+            if (ended)
+                continue;
+            bbs.insert(bb);
+            extra_use.insert(bb);
+            dominated.append(N->begin(), N->end());
         }
-        assert(dominated.empty());
     }
+
     // For each BB, find the first instruction(s) where the allocation is possibly dead.
     // If all successors are live, then there isn't one.
+    // If the BB has "invisible" uses, then there isn't one.
     // If all successors are dead, then it's the first instruction after the last use
     // within the BB.
     // If some successors are live and others are dead, it's the first instruction in

test/llvmpasses/alloc-opt-pass.ll

@@ -103,6 +103,7 @@ declare {}*** @julia.get_pgcstack()
 declare noalias {} addrspace(10)* @julia.gc_alloc_obj(i8*, i64, {} addrspace(10)*)
 declare {}* @julia.pointer_from_objref({} addrspace(11)*)
 declare void @llvm.memcpy.p11i8.p0i8.i64(i8 addrspace(11)* nocapture writeonly, i8* nocapture readonly, i64, i32, i1)
+declare void @llvm.memcpy.p0i8.p0i8.i64(i8* nocapture writeonly, i8* nocapture readonly, i64, i1)
 declare token @llvm.julia.gc_preserve_begin(...)
 declare void @llvm.julia.gc_preserve_end(token)
 
@@ -139,3 +140,25 @@ L2:
   ret void
 }
 ; CHECK-LABEL: }{{$}}
+
+; CHECK-LABEL: @lifetime_no_preserve_end
+; CHECK: alloca
+; CHECK-NOT: call token(...) @llvm.julia.gc_preserve_begin
+; CHECK: call void @llvm.lifetime.start
+; CHECK-NOT: call void @llvm.lifetime.end
+define void @lifetime_no_preserve_end({}* noalias nocapture noundef nonnull sret({}) %0) {
+  %pgcstack = call {}*** @julia.get_pgcstack()
+  %ptls = call {}*** @julia.ptls_states()
+  %ptls_i8 = bitcast {}*** %ptls to i8*
+  %v = call noalias {} addrspace(10)* @julia.gc_alloc_obj(i8* %ptls_i8, i64 8, {} addrspace(10)* @tag)
+  %token = call token (...) @llvm.julia.gc_preserve_begin({} addrspace(10)* %v)
+  %v_derived = addrspacecast {} addrspace(10)* %v to {} addrspace(11)*
+  %ptr = call nonnull {}* @julia.pointer_from_objref({} addrspace(11)* %v_derived)
+  %ptr_raw = bitcast {}* %ptr to i8*
+  call void @external_function() ; safepoint
+  %ret_raw = bitcast {}* %0 to i8*
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 8 %ret_raw, i8 * align 8 %ptr_raw, i64 0, i1 false)
+  %ret_raw2 = bitcast {}* %0 to i8*
+  ret void
+}
+; CHECK-LABEL: }{{$}}