AllocOpt: Fix stack lowering where alloca continas boxed and unboxed data (#55306)

wsmoses · vchuravy · giordano · wsmoses · commit fbc7d4f4c6d7 · 2024-08-06T11:03:10.000-04:00
Co-authored-by: Valentin Churavy &lt;vchuravy@users.noreply.github.com&gt;
Co-authored-by: Mosè Giordano &lt;giordano@users.noreply.github.com&gt;
Co-authored-by: Gabriel Baraldi &lt;baraldigabriel@gmail.com&gt;
diff --git a/src/llvm-alloc-helpers.cpp b/src/llvm-alloc-helpers.cpp
@@ -88,6 +88,8 @@ bool AllocUseInfo::addMemOp(Instruction *inst, unsigned opno, uint32_t offset,
     memop.isaggr = isa<StructType>(elty) || isa<ArrayType>(elty) || isa<VectorType>(elty);
     memop.isobjref = hasObjref(elty);
     auto &field = getField(offset, size, elty);
+    field.second.hasunboxed |= !hasObjref(elty) || (hasObjref(elty) && !isa<PointerType>(elty));
+
     if (field.second.hasobjref != memop.isobjref)
         field.second.multiloc = true; // can't split this field, since it contains a mix of references and bits
     if (!isstore)
@@ -189,6 +191,7 @@ void jl_alloc::runEscapeAnalysis(llvm::Instruction *I, EscapeAnalysisRequiredArg
                 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.has_unknown_unboxed |= !hasObjref(elty) || (hasObjref(elty) && !isa<PointerType>(elty));
                 required.use_info.hasunknownmem = true;
             } else if (!required.use_info.addMemOp(inst, 0, cur.offset,
                                                                inst->getType(),
@@ -280,6 +283,7 @@ void jl_alloc::runEscapeAnalysis(llvm::Instruction *I, EscapeAnalysisRequiredArg
                 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.has_unknown_unboxed |= !hasObjref(elty) || (hasObjref(elty) && !isa<PointerType>(elty));
                 required.use_info.hasunknownmem = true;
             } else if (!required.use_info.addMemOp(inst, use->getOperandNo(),
                                                                cur.offset, storev->getType(),
@@ -301,10 +305,14 @@ void jl_alloc::runEscapeAnalysis(llvm::Instruction *I, EscapeAnalysisRequiredArg
             }
             required.use_info.hasload = true;
             auto storev = isa<AtomicCmpXchgInst>(inst) ? cast<AtomicCmpXchgInst>(inst)->getNewValOperand() : cast<AtomicRMWInst>(inst)->getValOperand();
+            Type *elty = storev->getType();
             if (cur.offset == UINT32_MAX || !required.use_info.addMemOp(inst, use->getOperandNo(),
-                                                               cur.offset, storev->getType(),
+                                                               cur.offset, elty,
                                                                true, required.DL)) {
                 LLVM_DEBUG(dbgs() << "Atomic inst has unknown offset\n");
+                required.use_info.has_unknown_objref |= hasObjref(elty);
+                required.use_info.has_unknown_objrefaggr |= hasObjref(elty) && !isa<PointerType>(elty);
+                required.use_info.has_unknown_unboxed |= !hasObjref(elty) || (hasObjref(elty) && !isa<PointerType>(elty));
                 required.use_info.hasunknownmem = true;
             }
             required.use_info.refload = true;
diff --git a/src/llvm-alloc-helpers.h b/src/llvm-alloc-helpers.h
@@ -46,6 +46,8 @@ namespace jl_alloc {
         bool hasaggr:1;
         bool multiloc:1;
         bool hasload:1;
+        // The alloc has a unboxed object at this offset.
+        bool hasunboxed:1;
         llvm::Type *elty;
         llvm::SmallVector<MemOp,4> accesses;
         Field(uint32_t size, llvm::Type *elty)
@@ -54,6 +56,7 @@ namespace jl_alloc {
               hasaggr(false),
               multiloc(false),
               hasload(false),
+              hasunboxed(false),
               elty(elty)
         {
         }
@@ -93,6 +96,9 @@ namespace jl_alloc {
         // The alloc has an aggregate Julia object reference not in an explicit field.
         bool has_unknown_objrefaggr:1;
 
+        // The alloc has an unboxed object at an unknown offset.
+        bool has_unknown_unboxed:1;
+
         void reset()
         {
             escaped = false;
@@ -107,6 +113,7 @@ namespace jl_alloc {
             haserror = false;
             has_unknown_objref = false;
             has_unknown_objrefaggr = false;
+            has_unknown_unboxed = false;
             uses.clear();
             preserves.clear();
             memops.clear();
diff --git a/src/llvm-alloc-opt.cpp b/src/llvm-alloc-opt.cpp
@@ -252,10 +252,12 @@ void Optimizer::optimizeAll()
             removeAlloc(orig);
             continue;
         }
+        bool has_unboxed = use_info.has_unknown_unboxed;
         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;
+            has_unboxed |= field.hasunboxed;
             if (field.hasobjref) {
                 has_ref = true;
                 // This can be relaxed a little based on hasload
@@ -284,6 +286,19 @@ void Optimizer::optimizeAll()
             splitOnStack(orig);
             continue;
         }
+        // The move to stack code below, if has_ref is set, changes the allocation to an array of jlvalue_t's. This is fine
+        // if all objects are jlvalue_t's. However, if part of the allocation is an unboxed value (e.g. it is a { float, jlvaluet }),
+        // then moveToStack will create a [2 x jlvaluet] bitcast to { float, jlvaluet }.
+        // This later causes the GC rooting pass, to miss-characterize the float as a pointer to a GC value
+        if (has_unboxed && has_ref) {
+            REMARK([&]() {
+                return OptimizationRemarkMissed(DEBUG_TYPE, "Escaped", orig)
+                    << "GC allocation could not be split since it contains both boxed and unboxed values, unable to move to stack " << ore::NV("GC Allocation", orig);
+            });
+            if (use_info.hastypeof)
+                optimizeTag(orig);
+            continue;
+        }
         REMARK([&](){
             return OptimizationRemark(DEBUG_TYPE, "Stack Move Allocation", orig)
                 << "GC allocation moved to stack " << ore::NV("GC Allocation", orig);
diff --git a/test/llvmpasses/alloc-opt-bits.ll b/test/llvmpasses/alloc-opt-bits.ll
@@ -0,0 +1,41 @@
+; This file is a part of Julia. License is MIT: https://julialang.org/license
+
+
+; RUN: opt -enable-new-pm=0 --opaque-pointers=0 -load libjulia-codegen%shlibext -AllocOpt -LateLowerGCFrame -FinalLowerGC -S %s | FileCheck %s --check-prefixes=CHECK,TYPED
+; RUN: opt -enable-new-pm=1 --opaque-pointers=0 --load-pass-plugin=libjulia-codegen%shlibext -passes='function(AllocOpt,LateLowerGCFrame),FinalLowerGC' -S %s | FileCheck %s --check-prefixes=CHECK,TYPED
+
+; RUN: opt -enable-new-pm=0 --opaque-pointers=1 -load libjulia-codegen%shlibext -AllocOpt -LateLowerGCFrame -FinalLowerGC -S %s | FileCheck %s --check-prefixes=CHECK,OPAQUE
+; RUN: opt -enable-new-pm=1 --opaque-pointers=1 --load-pass-plugin=libjulia-codegen%shlibext -passes='function(AllocOpt,LateLowerGCFrame),FinalLowerGC' -S %s | FileCheck %s --check-prefixes=CHECK,OPAQUE
+
+
+@tag = external addrspace(10) global {}
+
+@glob = external addrspace(10) global {}
+
+; Test that the gc_preserve intrinsics are deleted directly.
+
+; CHECK-LABEL: @ptr_and_bits
+; CHECK-NOT: alloca 
+; OPAQUE: call noalias {} addrspace(10) @julia.gc_alloc_obj
+; TYPED: call noalias ptr addrspace(10) @julia.gc_alloc_obj
+
+define void @ptr_and_bits(i8* %fptr, i1 %b, i1 %b2, i32 %idx) {
+  %pgcstack = call {}*** @julia.get_pgcstack()
+  %gcstack = bitcast {}*** %pgcstack to {}**
+  %current_task = getelementptr inbounds {}*, {}** %gcstack, i64 -12
+  %v = call noalias {} addrspace(10)* @julia.gc_alloc_obj({}** %current_task, i64 16, {} addrspace(10) @tag)
+  %v2 = bitcast {} addrspace(10)* to { i64, {} addrspace(10)* } addrspace(10)*
+  %g0 = getelementptr { i64, {} addrspace(10)* }, { i64, {} addrspace(10)* }* %v2, i32 %idx, i32 1
+  store {} addrspace(10)* @glob, {} addrspace(10)* addrspace(10)* %g0
+  
+  %g1 = getelementptr { i64, {} addrspace(10)* }, { i64, {} addrspace(10)* }* %v2, i32 %idx, i32 0
+  store i64 7, i75 addrspace(10) %g1
+
+  %res = load {} addrspace(10)*, {} addrspace(10)* addrspace(10)* %g0
+  %res2 = load i64, i64 addrspace(10) %g1
+  ret void
+}
+
+declare noalias {} addrspace(10) @julia.gc_alloc_obj({}**, i64, {} addrspace(10))
+
+declare {}*** @julia.get_pgcstack()
