optimizer: fix alloc opt on unknown offset with references (#47076)

Fixes issued mentioned in #47075 (comment)

(cherry picked from commit a68235c)

Author: wsmoses

src/llvm-alloc-helpers.cpp

@@ -161,7 +161,12 @@ void jl_alloc::runEscapeAnalysis(llvm::Instruction *I, EscapeAnalysisRequiredArg
     auto check_inst = [&] (Instruction *inst, Use *use) {
         if (isa<LoadInst>(inst)) {
             required.use_info.hasload = true;
-            if (cur.offset == UINT32_MAX || !required.use_info.addMemOp(inst, 0, cur.offset,
+            if (cur.offset == UINT32_MAX) {
+                auto elty = inst->getType();
+                required.use_info.has_unknown_objref |= hasObjref(elty);
+                required.use_info.has_unknown_objrefaggr |= hasObjref(elty) && !isa<PointerType>(elty);
+                required.use_info.hasunknownmem = true;
+            } else if (!required.use_info.addMemOp(inst, 0, cur.offset,
                                                                inst->getType(),
                                                                false, required.DL))
                 required.use_info.hasunknownmem = true;
@@ -229,7 +234,12 @@ void jl_alloc::runEscapeAnalysis(llvm::Instruction *I, EscapeAnalysisRequiredArg
                 return false;
             }
             auto storev = store->getValueOperand();
-            if (cur.offset == UINT32_MAX || !required.use_info.addMemOp(inst, use->getOperandNo(),
+            if (cur.offset == UINT32_MAX) {
+                auto elty = storev->getType();
+                required.use_info.has_unknown_objref |= hasObjref(elty);
+                required.use_info.has_unknown_objrefaggr |= hasObjref(elty) && !isa<PointerType>(elty);
+                required.use_info.hasunknownmem = true;
+            } else if (!required.use_info.addMemOp(inst, use->getOperandNo(),
                                                                cur.offset, storev->getType(),
                                                                true, required.DL))
                 required.use_info.hasunknownmem = true;

src/llvm-alloc-helpers.h

@@ -87,6 +87,11 @@ namespace jl_alloc {
         // The object is used in an error function
         bool haserror:1;
 
+        // The alloc has a Julia object reference not in an explicit field.
+        bool has_unknown_objref:1;
+        // The alloc has an aggregate Julia object reference not in an explicit field.
+        bool has_unknown_objrefaggr:1;
+
         void reset()
         {
             escaped = false;
@@ -99,6 +104,8 @@ namespace jl_alloc {
             hasunknownmem = false;
             returned = false;
             haserror = false;
+            has_unknown_objref = false;
+            has_unknown_objrefaggr = false;
             uses.clear();
             preserves.clear();
             memops.clear();

src/llvm-alloc-opt.cpp

@@ -233,8 +233,8 @@ void Optimizer::optimizeAll()
             removeAlloc(orig);
             continue;
         }
-        bool has_ref = false;
-        bool has_refaggr = false;
+        bool has_ref = use_info.has_unknown_objref;
+        bool has_refaggr = use_info.has_unknown_objrefaggr;
         for (auto memop: use_info.memops) {
             auto &field = memop.second;
             if (field.hasobjref) {
@@ -577,7 +577,9 @@ void Optimizer::moveToStack(CallInst *orig_inst, size_t sz, bool has_ref)
         // treat this as a non-mem2reg'd alloca
         // The ccall root and GC preserve handling below makes sure that
         // the alloca isn't optimized out.
-        buff = prolog_builder.CreateAlloca(pass.T_prjlvalue);
+        const DataLayout &DL = F.getParent()->getDataLayout();
+        auto asize = ConstantInt::get(Type::getInt64Ty(prolog_builder.getContext()), sz / DL.getTypeAllocSize(pass.T_prjlvalue));
+        buff = prolog_builder.CreateAlloca(pass.T_prjlvalue, asize);
         buff->setAlignment(Align(align));
         ptr = cast<Instruction>(prolog_builder.CreateBitCast(buff, pass.T_pint8));
     }

test/llvmpasses/alloc-opt-unsized.ll

@@ -0,0 +1,35 @@
+; RUN: opt -enable-new-pm=0 -load libjulia-codegen%shlibext -AllocOpt -S %s | FileCheck %s
+
+source_filename = "text"
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128-ni:10:11:12:13"
+target triple = "x86_64-linux-gnu"
+
+declare {}*** @julia.get_pgcstack()
+
+declare {} addrspace(10)* @julia.gc_alloc_obj({}**, i64, {} addrspace(10)*)
+
+declare void @julia.write_barrier({} addrspace(10)*, ...)
+
+define void @diffejulia_objective__1864_inner_1wrap({} addrspace(10)* %arg, i64 %iv.i) {
+entry:
+  %i5 = call {}*** @julia.get_pgcstack()
+  %i13 = bitcast {}*** %i5 to {}**
+  %i14 = getelementptr inbounds {}*, {}** %i13, i64 -12
+  %i18 = call noalias nonnull dereferenceable(8000) dereferenceable_or_null(8000) {} addrspace(10)* @julia.gc_alloc_obj({}** %i14, i64 8000, {} addrspace(10)* addrspacecast ({}* inttoptr (i64 139756155247504 to {}*) to {} addrspace(10)*))
+  %_malloccache.i = bitcast {} addrspace(10)* %i18 to {} addrspace(10)* addrspace(10)*
+  %i23 = getelementptr inbounds {} addrspace(10)*, {} addrspace(10)* addrspace(10)* %_malloccache.i, i64 %iv.i
+  store {} addrspace(10)* %arg, {} addrspace(10)* addrspace(10)* %i23, align 8
+  %i24 = bitcast {} addrspace(10)* addrspace(10)* %_malloccache.i to {} addrspace(10)*
+  call void ({} addrspace(10)*, ...) @julia.write_barrier({} addrspace(10)* %i24, {} addrspace(10)* %arg)
+  %l = load {} addrspace(10)*, {} addrspace(10)* addrspace(10)* %i23
+  ret void
+}
+
+; CHECK:   %[[i0:.+]] = alloca {} addrspace(10)*, i64 1000, align 16
+; CHECK:   %[[i1:.+]] = bitcast {} addrspace(10)** %[[i0]] to i8*
+; CHECK:   %i18 = bitcast i8* %[[i1]] to {}*
+; CHECK:   %_malloccache.i = bitcast {}* %i18 to {} addrspace(10)**
+; CHECK:   %i23 = getelementptr inbounds {} addrspace(10)*, {} addrspace(10)** %_malloccache.i, i64 %iv.i
+; CHECK:   store {} addrspace(10)* %arg, {} addrspace(10)** %i23, align 8
+; CHECK:   %i24 = bitcast {} addrspace(10)** %_malloccache.i to {}*
+; CHECK:   %l = load {} addrspace(10)*, {} addrspace(10)** %i23, align 8