Merge pull request #50315 from JuliaLang/pc/fix-mssa-preserve

Fix MemorySSA preservation in julia-licm
JuliaLang · Jun 28, 2023 · 1177b54 · 1177b54
2 parents 5c070f4 + f5faa08
commit 1177b54
Show file tree

Hide file tree

Showing 3 changed files with 245 additions and 6 deletions.
diff --git a/src/llvm-julia-licm.cpp b/src/llvm-julia-licm.cpp
@@ -176,7 +176,7 @@ struct JuliaLICM : public JuliaPassContext {
  // Lazy initialization of exit blocks insertion points.
  bool exit_pts_init = false;
  SmallVector<Instruction*, 8> _exit_pts;
- auto get_exit_pts = [&] () -> ArrayRef<Instruction*> {
+ auto get_exit_pts = [&] () -> MutableArrayRef<Instruction*> {
  if (!exit_pts_init) {
  exit_pts_init = true;
  SmallVector<BasicBlock*, 8> exit_bbs;
@@ -242,6 +242,7 @@ struct JuliaLICM : public JuliaPassContext {
  }
  ++SunkPreserveEnd;
  moveInstructionBefore(*call, *exit_pts[0], MSSAU, SE, MemorySSA::Beginning);
+ exit_pts[0] = call;
  LLVM_DEBUG(dbgs() << "Sunk gc_preserve_end: " << *call << "\n");
  REMARK([&](){
  return OptimizationRemark(DEBUG_TYPE, "Sunk", call)
@@ -250,6 +251,7 @@ struct JuliaLICM : public JuliaPassContext {
  for (unsigned i = 1; i < exit_pts.size(); i++) {
  // Clone exit
  auto CI = CallInst::Create(call, {}, exit_pts[i]);
+ exit_pts[i] = CI;
  createNewInstruction(CI, call, MSSAU);
  LLVM_DEBUG(dbgs() << "Cloned and sunk gc_preserve_end: " << *CI << "\n");
  REMARK([&](){
@@ -345,11 +347,8 @@ struct JuliaLICM : public JuliaPassContext {
  auto align = Align(DL.getPointerSize(0));
  auto clear_obj = builder.CreateMemSet(obj_i8, ConstantInt::get(Type::getInt8Ty(call->getContext()), 0), call->getArgOperand(1), align);
  if (MSSAU.getMemorySSA()) {
- auto alloc_mdef = MSSAU.getMemorySSA()->getMemoryAccess(call);
- assert(isa<MemoryDef>(alloc_mdef) && "Expected alloc to be associated with a memory def!");
- auto clear_mdef = MSSAU.createMemoryAccessAfter(clear_obj, nullptr, alloc_mdef);
- assert(isa<MemoryDef>(clear_mdef) && "Expected memset to be associated with a memory def!");
- (void) clear_mdef;
+ auto clear_mdef = MSSAU.createMemoryAccessInBB(clear_obj, nullptr, clear_obj->getParent(), MemorySSA::BeforeTerminator);
+ MSSAU.insertDef(cast<MemoryDef>(clear_mdef), true);
  }
  changed = true;
  }

diff --git a/test/llvmpasses/julia-licm-memoryssa.ll b/test/llvmpasses/julia-licm-memoryssa.ll
@@ -0,0 +1,171 @@
+; COM: NewPM-only test, tests that memoryssa is preserved correctly
+
+; RUN: opt -enable-new-pm=1 --opaque-pointers=0 --load-pass-plugin=libjulia-codegen%shlibext -passes='function(loop-mssa(JuliaLICM),print<memoryssa>)' -S -o /dev/null %s 2>&1 | FileCheck %s --check-prefixes=CHECK,TYPED
+
+; RUN: opt -enable-new-pm=1 --opaque-pointers=1 --load-pass-plugin=libjulia-codegen%shlibext -passes='function(loop-mssa(JuliaLICM),print<memoryssa>)' -S -o /dev/null %s 2>&1 | FileCheck %s --check-prefixes=CHECK,OPAQUE
+
+@tag = external addrspace(10) global {}, align 16
+
+declare void @julia.write_barrier({} addrspace(10)*, ...)
+
+declare {}*** @julia.get_pgcstack()
+
+declare token @llvm.julia.gc_preserve_begin(...)
+
+declare void @llvm.julia.gc_preserve_end(token)
+
+declare void @mssa_use({} addrspace(10)*)
+
+declare noalias nonnull {} addrspace(10)* @julia.gc_alloc_obj({}**, i64, {} addrspace(10)*)
+
+; COM: check basic preserve hoist/sink functionality
+; CHECK-LABEL: MemorySSA for function: hoist_sink_preserves
+; CHECK-LABEL: @hoist_sink_preserves
+define void @hoist_sink_preserves({} addrspace(10)* %obj, i1 %ret) {
+; CHECK: top:
+top:
+; CHECK-NEXT: [[PGCSTACK:[0-9]+]] = MemoryDef(liveOnEntry)
+ %pgcstack = call {}*** @julia.get_pgcstack()
+ %current_task = bitcast {}*** %pgcstack to {}**
+; CHECK: br label %preheader
+ br label %preheader
+; CHECK: preheader:
+preheader:
+; CHECK-NEXT: [[PRESERVE_TOKEN:[0-9]+]] = MemoryDef([[PGCSTACK]])
+; CHECK-NEXT: %preserve_token = call token (...) @llvm.julia.gc_preserve_begin
+; CHECK-NEXT: br label %loop
+ br label %loop
+; CHECK: loop:
+loop:
+; CHECK-NOT: call token (...) @llvm.julia.gc_preserve_begin
+ %preserve_token = call token (...) @llvm.julia.gc_preserve_begin({} addrspace(10)* %obj)
+; CHECK-NOT: call void @llvm.julia.gc_preserve_end
+ call void @llvm.julia.gc_preserve_end(token %preserve_token)
+; CHECK-NEXT: [[MPHI:[0-9]+]] = MemoryPhi({preheader,[[PRESERVE_TOKEN]]},{loop,[[MPHI]]})
+; CHECK-NEXT: br i1 %ret
+ br i1 %ret, label %return, label %loop
+; CHECK: return:
+return:
+; CHECK-NEXT: [[PRESERVE_END:[0-9]+]] = MemoryDef([[MPHI]])
+; CHECK-NEXT: call void @llvm.julia.gc_preserve_end(token %preserve_token)
+; CHECK-NEXT: [[MSSA_USE:[0-9]+]] = MemoryDef([[PRESERVE_END]])
+; CHECK-NEXT: call void @mssa_use
+ call void @mssa_use({} addrspace(10)* %obj)
+; CHECK-NEXT: ret void
+ ret void
+}
+
+; COM: check sink functionality when there are multiple loop exit blocks
+; CHECK-LABEL: MemorySSA for function: hoist_multisink_preserves
+; CHECK-LABEL: @hoist_multisink_preserves
+define void @hoist_multisink_preserves({} addrspace(10)* %obj, i1 %ret) {
+; CHECK: top:
+top:
+; CHECK-NEXT: [[PGCSTACK:[0-9]+]] = MemoryDef(liveOnEntry)
+ %pgcstack = call {}*** @julia.get_pgcstack()
+ %current_task = bitcast {}*** %pgcstack to {}**
+; CHECK: br label %preheader
+ br label %preheader
+; CHECK: preheader:
+preheader:
+; CHECK-NEXT: [[PRESERVE_TOKEN:[0-9]+]] = MemoryDef([[PGCSTACK]])
+; CHECK-NEXT: %preserve_token = call token (...) @llvm.julia.gc_preserve_begin
+; CHECK-NEXT: br label %loop
+ br label %loop
+; CHECK: loop:
+loop:
+; CHECK-NOT: call token (...) @llvm.julia.gc_preserve_begin
+ %preserve_token = call token (...) @llvm.julia.gc_preserve_begin({} addrspace(10)* %obj)
+; CHECK-NOT: call void @llvm.julia.gc_preserve_end
+ call void @llvm.julia.gc_preserve_end(token %preserve_token)
+; CHECK-NEXT: [[MPHI:[0-9]+]] = MemoryPhi({preheader,[[PRESERVE_TOKEN]]},{loop2,[[MPHI]]})
+; CHECK-NEXT: br i1 %ret
+ br i1 %ret, label %return, label %loop2
+; CHECK: loop2:
+loop2:
+; CHECK-NEXT: br i1 %ret
+ br i1 %ret, label %return2, label %loop
+; CHECK: return:
+return:
+; CHECK-NEXT: [[PRESERVE_END_1:[0-9]+]] = MemoryDef([[MPHI]])
+; CHECK-NEXT: call void @llvm.julia.gc_preserve_end(token %preserve_token)
+; CHECK-NEXT: [[MSSA_USE:[0-9]+]] = MemoryDef([[PRESERVE_END_1]])
+; CHECK-NEXT: call void @mssa_use
+ call void @mssa_use({} addrspace(10)* %obj)
+; CHECK-NEXT: ret void
+ ret void
+; CHECK: return2:
+return2:
+; CHECK-NEXT: [[PRESERVE_END_2:[0-9]+]] = MemoryDef([[MPHI]])
+; CHECK-NEXT: call void @llvm.julia.gc_preserve_end(token %preserve_token)
+; CHECK-NEXT: [[MSSA_USE:[0-9]+]] = MemoryDef([[PRESERVE_END_2]])
+; CHECK-NEXT: call void @mssa_use
+ call void @mssa_use({} addrspace(10)* %obj)
+; CHECK-NEXT: ret void
+ ret void
+}
+
+define void @hoist_allocation({} addrspace(10)* %obj, i1 %ret) {
+; CHECK: top:
+top:
+; CHECK-NEXT: [[PGCSTACK:[0-9]+]] = MemoryDef(liveOnEntry)
+ %pgcstack = call {}*** @julia.get_pgcstack()
+ %current_task = bitcast {}*** %pgcstack to {}**
+ br label %preheader
+; CHECK: preheader:
+preheader:
+; CHECK-NEXT: [[ALLOC:[0-9]+]] = MemoryDef([[PGCSTACK]])
+
+; TYPED-NEXT: %alloc = call {} addrspace(10)* @julia.gc_alloc_obj({}** %current_task, i64 0, {} addrspace(10)* @tag)
+; TYPED-NEXT: %[[BCAST:.*]] = bitcast {} addrspace(10)* %alloc to i8 addrspace(10)*
+
+; OPAQUE-NEXT: %alloc = call ptr addrspace(10) @julia.gc_alloc_obj(ptr %current_task, i64 0, ptr addrspace(10) @tag)
+
+; CHECK-NEXT: [[MSET:[0-9]+]] = MemoryDef([[ALLOC]])
+; CHECK-NEXT: call void @llvm.memset
+; CHECK-NEXT: br label %loop
+ br label %loop
+; CHECK: loop:
+loop:
+; CHECK-NOT: %alloc
+; CHECK-NOT: @julia.gc_alloc_obj
+ %alloc = call {} addrspace(10)* @julia.gc_alloc_obj({}** %current_task, i64 0, {} addrspace(10)* @tag)
+; CHECK-NEXT: [[MPHI:[0-9]+]] = MemoryPhi({preheader,[[MSET]]},{loop,[[MPHI]]})
+ br i1 %ret, label %return, label %loop
+; CHECK: return:
+return:
+; CHECK-NEXT: [[MSSA_USE:[0-9]+]] = MemoryDef([[MPHI]])
+; CHECK-NEXT: call void @mssa_use
+ call void @mssa_use({} addrspace(10)* %obj)
+; CHECK-NEXT: ret void
+ ret void
+}
+
+define void @hoist_write_barrier({} addrspace(10)* %obj, i1 %ret) {
+; CHECK: top:
+top:
+; CHECK-NEXT: [[PGCSTACK:[0-9]+]] = MemoryDef(liveOnEntry)
+ %pgcstack = call {}*** @julia.get_pgcstack()
+ %current_task = bitcast {}*** %pgcstack to {}**
+ br label %preheader
+; CHECK: preheader:
+preheader:
+; CHECK-NEXT: [[WB:[0-9]+]] = MemoryDef([[PGCSTACK]])
+; CHECK-NEXT: call void
+; CHECK-SAME: @julia.write_barrier
+; CHECK-NEXT: br label %loop
+ br label %loop
+; CHECK: loop:
+loop:
+; CHECK-NOT: @julia.write_barrier
+ call void ({} addrspace(10)*, ...) @julia.write_barrier({} addrspace(10)* %obj)
+; CHECK-NEXT: [[MPHI:[0-9]+]] = MemoryPhi({preheader,[[WB]]},{loop,[[MPHI]]})
+ br i1 %ret, label %return, label %loop
+; CHECK: return:
+return:
+; CHECK-NEXT: [[MSSA_USE:[0-9]+]] = MemoryDef([[MPHI]])
+; CHECK-NEXT: call void @mssa_use
+ call void @mssa_use({} addrspace(10)* %obj)
+; CHECK-NEXT: ret void
+ ret void
+}
diff --git a/test/llvmpasses/julia-licm.ll b/test/llvmpasses/julia-licm.ll
@@ -12,6 +12,75 @@ declare void @julia.write_barrier({}*, ...)
 
 declare {}*** @julia.get_pgcstack()
 
+declare token @llvm.julia.gc_preserve_begin(...)
+
+declare void @llvm.julia.gc_preserve_end(token)
+
+; COM: check basic preserve hoist/sink functionality
+; CHECK-LABEL: @hoist_sink_preserves
+define void @hoist_sink_preserves({} addrspace(10)* %obj, i1 %ret) {
+top:
+ %pgcstack = call {}*** @julia.get_pgcstack()
+ %current_task = bitcast {}*** %pgcstack to {}**
+; CHECK: br label %preheader
+ br label %preheader
+; CHECK: preheader:
+preheader:
+; CHECK-NEXT: %preserve_token = call token (...) @llvm.julia.gc_preserve_begin
+; CHECK-NEXT: br label %loop
+ br label %loop
+; CHECK: loop:
+loop:
+; CHECK-NOT: call token (...) @llvm.julia.gc_preserve_begin
+ %preserve_token = call token (...) @llvm.julia.gc_preserve_begin({} addrspace(10)* %obj)
+; CHECK-NOT: call void @llvm.julia.gc_preserve_end
+ call void @llvm.julia.gc_preserve_end(token %preserve_token)
+; CHECK-NEXT: br i1 %ret
+ br i1 %ret, label %return, label %loop
+; CHECK: return:
+return:
+; CHECK-NEXT: call void @llvm.julia.gc_preserve_end(token %preserve_token)
+; CHECK-NEXT: ret void
+ ret void
+}
+
+; COM: check sink functionality when there are multiple loop exit blocks
+; CHECK-LABEL: @hoist_multisink_preserves
+define void @hoist_multisink_preserves({} addrspace(10)* %obj, i1 %ret) {
+top:
+ %pgcstack = call {}*** @julia.get_pgcstack()
+ %current_task = bitcast {}*** %pgcstack to {}**
+; CHECK: br label %preheader
+ br label %preheader
+; CHECK: preheader:
+preheader:
+; CHECK-NEXT: %preserve_token = call token (...) @llvm.julia.gc_preserve_begin
+; CHECK-NEXT: br label %loop
+ br label %loop
+; CHECK: loop:
+loop:
+; CHECK-NOT: call token (...) @llvm.julia.gc_preserve_begin
+ %preserve_token = call token (...) @llvm.julia.gc_preserve_begin({} addrspace(10)* %obj)
+; CHECK-NOT: call void @llvm.julia.gc_preserve_end
+ call void @llvm.julia.gc_preserve_end(token %preserve_token)
+; CHECK-NEXT: br i1 %ret
+ br i1 %ret, label %return, label %loop2
+; CHECK: loop2:
+loop2:
+; CHECK-NEXT: br i1 %ret
+ br i1 %ret, label %return2, label %loop
+; CHECK: return:
+return:
+; CHECK-NEXT: call void @llvm.julia.gc_preserve_end(token %preserve_token)
+; CHECK-NEXT: ret void
+ ret void
+; CHECK: return2:
+return2:
+; CHECK-NEXT: call void @llvm.julia.gc_preserve_end(token %preserve_token)
+; CHECK-NEXT: ret void
+ ret void
+}
+
 ; COM: check basic allocation hoisting functionality
 ; CHECK-LABEL: @julia_allocation_hoist
 define nonnull {} addrspace(10)* @julia_allocation_hoist(i64 signext %0) #0 {