Create tbaa_gcframe and decorate all stores and loads from the GC frame with it

Close #13301

Author: yuyichao

src/codegen.cpp

@@ -262,6 +262,7 @@ static Type *T_pppint8;
 static Type *T_void;
 
 // type-based alias analysis nodes.  Indentation of comments indicates hierarchy.
+static MDNode *tbaa_gcframe;        // GC frame
 static MDNode *tbaa_user;           // User data that is mutable
 static MDNode *tbaa_immut;          // User data inside a heap-allocated immutable
 static MDNode *tbaa_value;          // Julia value
@@ -3503,7 +3504,8 @@ static void allocate_gc_frame(BasicBlock *b0, jl_codectx_t *ctx)
 #endif
 }
 
-void jl_codegen_finalize_temp_arg(CallInst *ptlsStates, Type *T_pjlvalue);
+void jl_codegen_finalize_temp_arg(CallInst *ptlsStates, Type *T_pjlvalue,
+                                  MDNode *tbaa_gcframe);
 static void finalize_gc_frame(Function *F)
 {
     Module *M = F->getParent();
@@ -3525,7 +3527,7 @@ static void finalize_gc_frame(Function *F)
     if (!ptlsStates)
         return;
 
-    jl_codegen_finalize_temp_arg(ptlsStates, T_pjlvalue);
+    jl_codegen_finalize_temp_arg(ptlsStates, T_pjlvalue, tbaa_gcframe);
 
 #ifdef JULIA_ENABLE_THREADING
     if (imaging_mode) {
@@ -4891,6 +4893,7 @@ extern "C" void jl_fptr_to_llvm(jl_fptr_t fptr, jl_lambda_info_t *lam, int specs
 static void init_julia_llvm_env(Module *m)
 {
     MDNode *tbaa_root = mbuilder->createTBAARoot("jtbaa");
+    tbaa_gcframe = tbaa_make_child("jtbaa_gcframe",tbaa_root);
     tbaa_user = tbaa_make_child("jtbaa_user",tbaa_root);
     tbaa_value = tbaa_make_child("jtbaa_value",tbaa_root);
     tbaa_immut = tbaa_make_child("jtbaa_immut",tbaa_root);

src/llvm-gcroot.cpp

@@ -15,6 +15,7 @@
 
 #include <vector>
 #include <queue>
+#include <set>
 
 #include "julia.h"
 
@@ -36,7 +37,7 @@ static struct {
 
 class JuliaGCAllocator {
 public:
-    JuliaGCAllocator(CallInst *ptlsStates, Type *T_pjlvalue) :
+    JuliaGCAllocator(CallInst *ptlsStates, Type *T_pjlvalue, MDNode *tbaa) :
         F(*ptlsStates->getParent()->getParent()),
         M(*F.getParent()),
         T_int32(Type::getInt32Ty(F.getContext())),
@@ -47,7 +48,8 @@ class JuliaGCAllocator {
         gcroot_func(M.getFunction("julia.gc_root_decl")),
         gckill_func(M.getFunction("julia.gc_root_kill")),
         gc_store_func(M.getFunction("julia.gc_store")),
-        jlcall_frame_func(M.getFunction("julia.jlcall_frame_decl"))
+        jlcall_frame_func(M.getFunction("julia.jlcall_frame_decl")),
+        tbaa_gcframe(tbaa)
     {
 /* Algorithm sketch:
  *  Compute liveness for each basic block
@@ -74,6 +76,7 @@ Function *const gcroot_func;
 Function *const gckill_func;
 Function *const gc_store_func;
 Function *const jlcall_frame_func;
+MDNode *tbaa_gcframe;
 
 typedef std::pair<CallInst*, unsigned> frame_register;
 class liveness {
@@ -90,6 +93,42 @@ class liveness {
     };
 };
 
+    void tbaa_decorate_gcframe(Instruction *inst, std::set<Instruction*> &visited)
+    {
+        if (visited.find(inst) != visited.end())
+            return;
+        visited.insert(inst);
+#ifdef LLVM35
+        Value::user_iterator I = inst->user_begin(), E = inst->user_end();
+#else
+        Value::use_iterator I = inst->use_begin(), E = inst->use_end();
+#endif
+        for (;I != E;++I) {
+            Instruction *user = dyn_cast<Instruction>(*I);
+            if (!user) {
+                continue;
+            } else if (isa<GetElementPtrInst>(user)) {
+                if (__likely(user->getOperand(0) == inst)) {
+                    tbaa_decorate_gcframe(user, visited);
+                }
+            } else if (isa<StoreInst>(user)) {
+                if (user->getOperand(1) == inst) {
+                    user->setMetadata(llvm::LLVMContext::MD_tbaa, tbaa_gcframe);
+                }
+            } else if (isa<LoadInst>(user)) {
+                user->setMetadata(llvm::LLVMContext::MD_tbaa, tbaa_gcframe);
+            } else if (isa<BitCastInst>(user)) {
+                tbaa_decorate_gcframe(user, visited);
+            }
+        }
+    }
+
+    void tbaa_decorate_gcframe(Instruction *inst)
+    {
+        std::set<Instruction*> visited;
+        tbaa_decorate_gcframe(inst, visited);
+    }
+
 #ifndef NDEBUG // llvm assertions build
 // gdb debugging code for inspecting the bb_uses map
 void jl_dump_bb_uses(std::map<BasicBlock*, std::map<frame_register, liveness::id> > &bb_uses)
@@ -655,6 +694,7 @@ void allocate_frame()
         // finalize all of the jlcall frames by replacing all of the frames with the appropriate gep(tempslot)
         for (std::map<CallInst*, unsigned>::iterator frame = frame_offsets.begin(), framee = frame_offsets.end(); frame != framee; ++frame) {
             CallInst *gcroot = frame->first;
+            tbaa_decorate_gcframe(gcroot);
             Value* offset[1] = {ConstantInt::get(T_int32, frame->second)};
             GetElementPtrInst *gep = GetElementPtrInst::Create(LLVM37_param(NULL) tempSlot, makeArrayRef(offset));
             gep->insertAfter(last_gcframe_inst);
@@ -714,11 +754,13 @@ void allocate_frame()
                     }
                 }
 #endif
+                tbaa_decorate_gcframe(callInst);
                 callInst->replaceAllUsesWith(argTempi);
                 argTempi->takeName(callInst);
                 callInst->eraseFromParent();
                 // Initialize the slots for function variables to NULL
                 StoreInst *store = new StoreInst(V_null, argTempi);
+                store->setMetadata(llvm::LLVMContext::MD_tbaa, tbaa_gcframe);
                 store->insertAfter(argTempi);
                 last_gcframe_inst = store;
             }
@@ -744,6 +786,7 @@ void allocate_frame()
             Instruction *argTempi = GetElementPtrInst::Create(LLVM37_param(NULL) tempSlot, ArrayRef<Value*>(ConstantInt::get(T_int32, i)));
             argTempi->insertAfter(last_gcframe_inst);
             StoreInst *store = new StoreInst(V_null, argTempi);
+            store->setMetadata(llvm::LLVMContext::MD_tbaa, tbaa_gcframe);
             store->insertAfter(argTempi);
             last_gcframe_inst = store;
         }
@@ -763,10 +806,13 @@ void allocate_frame()
         DebugLoc noDbg;
         builder.SetCurrentDebugLocation(noDbg);
 
-        builder.CreateStore(ConstantInt::get(T_size, (argSpaceSize + maxDepth) << 1),
-                            builder.CreateBitCast(builder.CreateConstGEP1_32(gcframe, 0), T_size->getPointerTo()));
-        builder.CreateStore(builder.CreateLoad(builder.Insert(get_pgcstack(ptlsStates))),
-                            builder.CreatePointerCast(builder.CreateConstGEP1_32(gcframe, 1), PointerType::get(T_ppjlvalue,0)));
+        Instruction *inst =
+            builder.CreateStore(ConstantInt::get(T_size, (argSpaceSize + maxDepth) << 1),
+                                builder.CreateBitCast(builder.CreateConstGEP1_32(gcframe, 0), T_size->getPointerTo()));
+        inst->setMetadata(llvm::LLVMContext::MD_tbaa, tbaa_gcframe);
+        inst = builder.CreateStore(builder.CreateLoad(builder.Insert(get_pgcstack(ptlsStates))),
+                                   builder.CreatePointerCast(builder.CreateConstGEP1_32(gcframe, 1), PointerType::get(T_ppjlvalue,0)));
+        inst->setMetadata(llvm::LLVMContext::MD_tbaa, tbaa_gcframe);
         builder.CreateStore(gcframe, builder.Insert(get_pgcstack(ptlsStates)));
 
         // Finish by emitting the gc pops before any return
@@ -775,9 +821,11 @@ void allocate_frame()
                 builder.SetInsertPoint(I->getTerminator()); // set insert *before* Ret
                 Instruction *gcpop =
                     (Instruction*)builder.CreateConstGEP1_32(gcframe, 1);
-                builder.CreateStore(builder.CreatePointerCast(builder.CreateLoad(gcpop),
-                                                          T_ppjlvalue),
-                                    builder.Insert(get_pgcstack(ptlsStates)));
+                inst = builder.CreateLoad(gcpop);
+                inst->setMetadata(llvm::LLVMContext::MD_tbaa, tbaa_gcframe);
+                inst = builder.CreateStore(builder.CreatePointerCast(inst, T_ppjlvalue),
+                                           builder.Insert(get_pgcstack(ptlsStates)));
+                inst->setMetadata(llvm::LLVMContext::MD_tbaa, tbaa_gcframe);
             }
         }
     }
@@ -791,8 +839,9 @@ void allocate_frame()
 
 };
 
-void jl_codegen_finalize_temp_arg(CallInst *ptlsStates, Type *T_pjlvalue)
+void jl_codegen_finalize_temp_arg(CallInst *ptlsStates, Type *T_pjlvalue,
+                                  MDNode *tbaa)
 {
-    JuliaGCAllocator allocator(ptlsStates, T_pjlvalue);
+    JuliaGCAllocator allocator(ptlsStates, T_pjlvalue, tbaa);
     allocator.allocate_frame();
 }