[KeyInstr][JumpThreading] Remap atoms in blocks duplicated for threading

Author: OCHyams

llvm/lib/Transforms/Scalar/JumpThreading.cpp

@@ -2011,6 +2011,14 @@ void JumpThreadingPass::updateSSA(BasicBlock *BB, BasicBlock *NewBB,
   }
 }
 
+static void remapSourceAtoms(ValueToValueMapTy &VM, BasicBlock::iterator Begin,
+                             BasicBlock::iterator End) {
+  if (VM.AtomMap.empty())
+    return;
+  for (auto It = Begin; It != End; ++It)
+    RemapSourceAtom(&*It, VM);
+}
+
 /// Clone instructions in range [BI, BE) to NewBB.  For PHI nodes, we only clone
 /// arguments that come from PredBB.  Return the map from the variables in the
 /// source basic block to the variables in the newly created basic block.
@@ -2075,6 +2083,8 @@ void JumpThreadingPass::cloneInstructions(ValueToValueMapTy &ValueMapping,
     PHINode *NewPN = PHINode::Create(PN->getType(), 1, PN->getName(), NewBB);
     NewPN->addIncoming(PN->getIncomingValueForBlock(PredBB), PredBB);
     ValueMapping[PN] = NewPN;
+    if (const DebugLoc &DL = PN->getDebugLoc())
+      mapAtomInstance(DL, ValueMapping);
   }
 
   // Clone noalias scope declarations in the threaded block. When threading a
@@ -2103,8 +2113,10 @@ void JumpThreadingPass::cloneInstructions(ValueToValueMapTy &ValueMapping,
     adaptNoAliasScopes(New, ClonedScopes, Context);
 
     CloneAndRemapDbgInfo(New, &*BI);
+    if (const DebugLoc &DL = New->getDebugLoc())
+      mapAtomInstance(DL, ValueMapping);
 
-    if (RetargetDbgValueIfPossible(New))
+      if (RetargetDbgValueIfPossible(New))
       continue;
 
     // Remap operands to patch up intra-block references.
@@ -2330,6 +2342,7 @@ void JumpThreadingPass::threadThroughTwoBasicBlocks(BasicBlock *PredPredBB,
        {DominatorTree::Insert, PredPredBB, NewBB},
        {DominatorTree::Delete, PredPredBB, PredBB}});
 
+  remapSourceAtoms(ValueMapping, NewBB->begin(), NewBB->end());
   updateSSA(PredBB, NewBB, ValueMapping);
 
   // Clean up things like PHI nodes with single operands, dead instructions,

llvm/test/DebugInfo/KeyInstructions/Generic/jump-threading-2-bbs.ll

@@ -0,0 +1,68 @@
+; RUN: opt -S -passes=jump-threading,verify %s | FileCheck %s
+
+;; Modified from llvm/test/Transforms/JumpThreading/thread-two-bbs.ll
+;;
+;; JumpThreading duplicates bb.cond2 to thread through bb.file to bb.f2 or exit.
+;;
+;; Check the duplicated instructions get remapped atom groups.
+
+; CHECK: bb.cond2:
+; CHECK-NEXT: call void @f1()
+; CHECK-NEXT: %tobool1 = icmp eq i32 %cond2, 0, !dbg [[G1R2:!.*]]
+; CHECK-NEXT: br i1 %tobool1, label %exit, label %exit, !dbg [[G1R1:!.*]]
+
+; CHECK: bb.cond2.thread:
+; CHECK-NEXT: %tobool12 = icmp eq i32 %cond2, 0, !dbg [[G2R2:!.*]]
+; CHECK-NEXT: br i1 %tobool12, label %bb.f2, label %exit, !dbg [[G2R1:!.*]]
+
+; CHECK: [[G1R2]] = !DILocation(line: 1, column: 1, scope: ![[#]], atomGroup: 1, atomRank: 2)
+; CHECK: [[G1R1]] = !DILocation(line: 1, column: 1, scope: ![[#]], atomGroup: 1, atomRank: 1)
+; CHECK: [[G2R2]] = !DILocation(line: 1, column: 1, scope: ![[#]], atomGroup: 2, atomRank: 2)
+; CHECK: [[G2R1]] = !DILocation(line: 1, column: 1, scope: ![[#]], atomGroup: 2, atomRank: 1)
+
+@a = global i32 0, align 4
+
+define void @foo(i32 %cond1, i32 %cond2) !dbg !5 {
+entry:
+  %tobool = icmp eq i32 %cond1, 0
+  br i1 %tobool, label %bb.cond2, label %bb.f1
+
+bb.f1:                                            ; preds = %entry
+  call void @f1()
+  br label %bb.cond2
+
+bb.cond2:                                         ; preds = %bb.f1, %entry
+  %ptr = phi ptr [ null, %bb.f1 ], [ @a, %entry ]
+  %tobool1 = icmp eq i32 %cond2, 0, !dbg !9
+  br i1 %tobool1, label %bb.file, label %exit, !dbg !10
+
+bb.file:                                          ; preds = %bb.cond2
+  %cmp = icmp eq ptr %ptr, null
+  br i1 %cmp, label %exit, label %bb.f2
+
+bb.f2:                                            ; preds = %bb.file
+  call void @f2()
+  br label %exit
+
+exit:                                             ; preds = %bb.f2, %bb.file, %bb.cond2
+  ret void
+}
+
+declare void @f1()
+
+declare void @f2()
+
+!llvm.dbg.cu = !{!0}
+!llvm.debugify = !{!2, !3}
+!llvm.module.flags = !{!4}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C, file: !1, producer: "debugify", isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug)
+!1 = !DIFile(filename: "<stdin>", directory: "/")
+!2 = !{i32 16}
+!3 = !{i32 0}
+!4 = !{i32 2, !"Debug Info Version", i32 3}
+!5 = distinct !DISubprogram(name: "foo", linkageName: "foo", scope: null, file: !1, line: 1, type: !6, scopeLine: 1, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
+!6 = !DISubroutineType(types: !7)
+!7 = !{}
+!9 = !DILocation(line: 1, column: 1, scope: !5, atomGroup: 1, atomRank: 2)
+!10 = !DILocation(line: 1, column: 1, scope: !5, atomGroup: 1, atomRank: 1)