Merge pull request #23783 from JuliaLang/yyc/codegen/alloc

Fix support of gc_preserve intrinsics in the allocation optimization pass
JuliaLang · Sep 22, 2017 · 9036d30 · 9036d30
2 parents 72eed90 + fa925a8
commit 9036d30
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 11 deletions.
diff --git a/src/llvm-alloc-opt.cpp b/src/llvm-alloc-opt.cpp
@@ -62,6 +62,8 @@ static bool isBundleOperand(CallInst *call, unsigned idx)
  *
  * * load
  * * `pointer_from_objref`
+ * * Any real llvm intrinsics
+ * * gc preserve intrinsics
  * * `ccall` gcroot array (`jl_roots` operand bundle)
  * * store (as address)
  * * addrspacecast, bitcast, getelementptr
@@ -88,6 +90,7 @@ struct AllocOpt : public FunctionPass {
     Function *ptr_from_objref;
     Function *lifetime_start;
     Function *lifetime_end;
+    Function *gc_preserve_begin;
 
     Type *T_int8;
     Type *T_int32;
@@ -150,7 +153,8 @@ struct AllocOpt : public FunctionPass {
     bool checkInst(Instruction *I, CheckInstStack &stack, std::set<Instruction*> &uses,
                    bool &ignore_tag);
     void replaceUsesWith(Instruction *orig_i, Instruction *new_i, ReplaceUsesStack &stack);
-    void replaceIntrinsicUseWith(IntrinsicInst *call, Instruction *orig_i, Instruction *new_i);
+    void replaceIntrinsicUseWith(IntrinsicInst *call, Intrinsic::ID ID, Instruction *orig_i,
+                                 Instruction *new_i);
     bool isSafepoint(Instruction *inst);
     void getAnalysisUsage(AnalysisUsage &AU) const override
     {
@@ -333,6 +337,7 @@ bool AllocOpt::doInitialization(Module &M)
         return false;
 
     ptr_from_objref = M.getFunction("julia.pointer_from_objref");
+    gc_preserve_begin = M.getFunction("llvm.julia.gc_preserve_begin");
 
     T_prjlvalue = alloc_obj->getReturnType();
     T_pjlvalue = PointerType::get(cast<PointerType>(T_prjlvalue)->getElementType(), 0);
@@ -384,9 +389,16 @@ bool AllocOpt::checkInst(Instruction *I, CheckInstStack &stack, std::set<Instruc
         if (auto call = dyn_cast<CallInst>(inst)) {
             // TODO handle `memcmp`
             // None of the intrinsics should care if the memory is stack or heap allocated.
-            if (isa<IntrinsicInst>(call))
-                return true;
-            if (ptr_from_objref && ptr_from_objref == call->getCalledFunction())
+            auto callee = call->getCalledFunction();
+            if (auto II = dyn_cast<IntrinsicInst>(call)) {
+                if (II->getIntrinsicID()) {
+                    return true;
+                }
+                if (gc_preserve_begin && gc_preserve_begin == callee) {
+                    return true;
+                }
+            }
+            if (ptr_from_objref && ptr_from_objref == callee)
                 return true;
             auto opno = use->getOperandNo();
             // Uses in `jl_roots` operand bundle are not counted as escaping, everything else is.
@@ -445,11 +457,9 @@ bool AllocOpt::checkInst(Instruction *I, CheckInstStack &stack, std::set<Instruc
     }
 }
 
-void AllocOpt::replaceIntrinsicUseWith(IntrinsicInst *call, Instruction *orig_i,
-                                       Instruction *new_i)
+void AllocOpt::replaceIntrinsicUseWith(IntrinsicInst *call, Intrinsic::ID ID,
+                                       Instruction *orig_i, Instruction *new_i)
 {
-    Intrinsic::ID ID = call->getIntrinsicID();
-    assert(ID);
     auto nargs = call->getNumArgOperands();
     SmallVector<Value*, 8> args(nargs);
     SmallVector<Type*, 8> argTys(nargs);
@@ -549,10 +559,12 @@ void AllocOpt::replaceUsesWith(Instruction *orig_inst, Instruction *new_inst,
                 return;
             }
             if (auto intrinsic = dyn_cast<IntrinsicInst>(call)) {
-                replaceIntrinsicUseWith(intrinsic, orig_i, new_i);
-                return;
+                if (Intrinsic::ID ID = intrinsic->getIntrinsicID()) {
+                    replaceIntrinsicUseWith(intrinsic, ID, orig_i, new_i);
+                    return;
+                }
             }
-            // remove from operand bundle
+            // remove from operand bundle or arguments for gc_perserve_begin
             Type *new_t = new_i->getType();
             user->replaceUsesOfWith(orig_i, ConstantPointerNull::get(cast<PointerType>(new_t)));
         }

diff --git a/src/llvm-late-gc-lowering.cpp b/src/llvm-late-gc-lowering.cpp
@@ -764,6 +764,8 @@ State LateLowerGCFrame::LocalScan(Function &F) {
                         std::vector<int> args;
                         for (Use &U : CI->arg_operands()) {
                             Value *V = U;
+                            if (isa<Constant>(V))
+                                continue;
                             int Num = Number(S, V);
                             if (Num >= 0)
                                 args.push_back(Num);

diff --git a/test/codegen.jl b/test/codegen.jl
@@ -188,6 +188,13 @@ function two_breakpoint(a::Float64)
     ccall(:jl_breakpoint, Void, (Ref{Float64},), a)
 end
 
+function load_dummy_ref(x::Int)
+    r = Ref{Int}(x)
+    Base.@gc_preserve r begin
+        unsafe_load(Ptr{Int}(pointer_from_objref(r)))
+    end
+end
+
 if opt_level > 0
     breakpoint_f64_ir = get_llvm((a)->ccall(:jl_breakpoint, Void, (Ref{Float64},), a),
                                  Tuple{Float64})
@@ -198,6 +205,12 @@ if opt_level > 0
     two_breakpoint_ir = get_llvm(two_breakpoint, Tuple{Float64})
     @test !contains(two_breakpoint_ir, "jl_gc_pool_alloc")
     @test contains(two_breakpoint_ir, "llvm.lifetime.end")
+
+    @test load_dummy_ref(1234) === 1234
+    load_dummy_ref_ir = get_llvm(load_dummy_ref, Tuple{Int})
+    @test !contains(load_dummy_ref_ir, "jl_gc_pool_alloc")
+    # Hopefully this is reliable enough. LLVM should be able to optimize this to a direct return.
+    @test contains(load_dummy_ref_ir, "ret $Iptr %0")
 end
 
 # Issue 22770

diff --git a/test/llvmpasses/alloc-opt.jl b/test/llvmpasses/alloc-opt.jl
@@ -205,6 +205,25 @@ top:
 """)
 # CHECK-LABEL: }
 
+# CHECK-LABEL: @preserve_opt
+# CHECK: alloca i128, align 16
+# CHECK: call %jl_value_t*** @jl_get_ptls_states()
+# CHECK-NOT: @julia.gc_alloc_obj
+# CHECK-NOT: @jl_gc_pool_alloc
+println("""
+define void @preserve_opt(i8* %v22) {
+top:
+  %v6 = call %jl_value_t*** @jl_get_ptls_states()
+  %v18 = bitcast %jl_value_t*** %v6 to i8*
+  %v19 = call noalias %jl_value_t addrspace(10)* @julia.gc_alloc_obj(i8* %v18, $isz 16, %jl_value_t addrspace(10)* @tag)
+  %v20 = bitcast %jl_value_t addrspace(10)* %v19 to i8 addrspace(10)*
+  %v21 = addrspacecast i8 addrspace(10)* %v20 to i8 addrspace(11)*
+  %tok = call token (...) @llvm.julia.gc_preserve_begin(%jl_value_t addrspace(10)* %v19)
+  ret void
+}
+""")
+# CHECK-LABEL: }
+
 # CHECK: declare noalias %jl_value_t addrspace(10)* @jl_gc_pool_alloc(i8*,
 # CHECK: declare noalias %jl_value_t addrspace(10)* @jl_gc_big_alloc(i8*,
 println("""
@@ -213,6 +232,7 @@ declare %jl_value_t*** @jl_get_ptls_states()
 declare noalias %jl_value_t addrspace(10)* @julia.gc_alloc_obj(i8*, $isz, %jl_value_t addrspace(10)*)
 declare i64 @julia.pointer_from_objref(%jl_value_t addrspace(11)*)
 declare void @llvm.memcpy.p11i8.p0i8.i64(i8 addrspace(11)* nocapture writeonly, i8* nocapture readonly, i64, i32, i1)
+declare token @llvm.julia.gc_preserve_begin(...)
 
 !0 = !{!1, !1, i64 0}
 !1 = !{!"jtbaa_tag", !2, i64 0}