[KeyInstr][SimplifyCFG] Remap atoms when folding br to common succ into pred

SimplifyCFG folds `b` into `a`.

            +-----------+
            |           v
  --> a --> b --> c --> d -->
      |                 ^
      +-----------------+

Remap source atoms in `b` so that the duplicated instructions are analysed
independently to determine is_stmt positions. This is necessary as the
contents of `b` may be folded into multiple preds in this way.

Add multi-pred test.

Author: OCHyams

llvm/include/llvm/IR/DebugLoc.h

@@ -76,6 +76,16 @@ namespace llvm {
                                     LLVMContext &Ctx,
                                     DenseMap<const MDNode *, MDNode *> &Cache);
 
+    /// Return true if the source locations match, ignoring isImplicitCode and
+    /// source atom info.
+    bool isSameSourceLocation(const DebugLoc &Other) const {
+      if (get() == Other.get())
+        return true;
+      return ((bool)*this == (bool)Other) && getLine() == Other.getLine() &&
+             getCol() == Other.getCol() && getScope() == Other.getScope() &&
+             getInlinedAt() == Other.getInlinedAt();
+    }
+
     unsigned getLine() const;
     unsigned getCol() const;
     MDNode *getScope() const;

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -73,6 +73,7 @@
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/LockstepReverseIterator.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
@@ -1129,13 +1130,17 @@ static void cloneInstructionsIntoPredecessorBlockAndUpdateSSAUses(
 
     Instruction *NewBonusInst = BonusInst.clone();
 
-    if (!isa<DbgInfoIntrinsic>(BonusInst) &&
-        PTI->getDebugLoc() != NewBonusInst->getDebugLoc()) {
-      // Unless the instruction has the same !dbg location as the original
-      // branch, drop it. When we fold the bonus instructions we want to make
-      // sure we reset their debug locations in order to avoid stepping on
-      // dead code caused by folding dead branches.
-      NewBonusInst->setDebugLoc(DebugLoc());
+    if (!isa<DbgInfoIntrinsic>(BonusInst)) {
+      if (!NewBonusInst->getDebugLoc().isSameSourceLocation(
+              PTI->getDebugLoc())) {
+        // Unless the instruction has the same !dbg location as the original
+        // branch, drop it. When we fold the bonus instructions we want to make
+        // sure we reset their debug locations in order to avoid stepping on
+        // dead code caused by folding dead branches.
+        NewBonusInst->setDebugLoc(DebugLoc());
+      } else if (const DebugLoc &DL = NewBonusInst->getDebugLoc()) {
+        mapAtomInstance(DL, VMap);
+      }
     }
 
     RemapInstruction(NewBonusInst, VMap,
@@ -1182,6 +1187,19 @@ static void cloneInstructionsIntoPredecessorBlockAndUpdateSSAUses(
       U.set(NewBonusInst);
     }
   }
+
+  // Key Instructions: We may have propagated atom info into the pred. If the
+  // pred's terminator already has atom info do nothing as merging would drop
+  // one atom group anyway. If it doesn't, propagte the remapped atom group
+  // from BB's terminator.
+  if (auto &PredDL = PredBlock->getTerminator()->getDebugLoc()) {
+    auto &DL = BB->getTerminator()->getDebugLoc();
+    if (!PredDL->getAtomGroup() && DL && DL->getAtomGroup() &&
+        PredDL.isSameSourceLocation(DL)) {
+      PredBlock->getTerminator()->setDebugLoc(DL);
+      RemapSourceAtom(PredBlock->getTerminator(), VMap);
+    }
+  }
 }
 
 bool SimplifyCFGOpt::performValueComparisonIntoPredecessorFolding(

llvm/test/DebugInfo/KeyInstructions/Generic/simplifycfg-branch-fold.ll

@@ -0,0 +1,81 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt %s -S -passes=simplifycfg -bonus-inst-threshold=2 | FileCheck %s
+
+;; Block d gets folded into preds b and c. Check the cloned instructions get
+;; remapped DILocation atomGroup numbers in each of the preds. Additionally
+;; check that the branches each inherit the atomGroup of the folded branch.
+
+declare i32 @g(...)
+define void @f(i1 %c0, i1 %c1, i1 %c2, i32 %x, i32 %y) !dbg !17 {
+; CHECK-LABEL: define void @f(
+; CHECK-SAME: i1 [[C0:%.*]], i1 [[C1:%.*]], i1 [[C2:%.*]], i32 [[X:%.*]], i32 [[Y:%.*]]) !dbg [[DBG6:![0-9]+]] {
+; CHECK-NEXT:  [[A:.*:]]
+; CHECK-NEXT:    br i1 [[C0]], label %[[B:.*]], label %[[C:.*]]
+; CHECK:       [[B]]:
+; CHECK-NEXT:    [[AND_OLD:%.*]] = and i32 [[X]], [[Y]], !dbg [[DBG8:![0-9]+]]
+; CHECK-NEXT:    [[CMP_OLD:%.*]] = icmp eq i32 [[AND_OLD]], 0, !dbg [[DBG9:![0-9]+]]
+; CHECK-NEXT:    [[OR_COND1:%.*]] = select i1 [[C1]], i1 true, i1 [[CMP_OLD]], !dbg [[DBG10:![0-9]+]]
+; CHECK-NEXT:    br i1 [[OR_COND1]], label %[[F:.*]], label %[[E:.*]], !dbg [[DBG11:![0-9]+]]
+; CHECK:       [[C]]:
+; CHECK-NEXT:    [[AND:%.*]] = and i32 [[X]], [[Y]], !dbg [[DBG12:![0-9]+]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[AND]], 0, !dbg [[DBG13:![0-9]+]]
+; CHECK-NEXT:    [[OR_COND:%.*]] = select i1 [[C2]], i1 true, i1 [[CMP]], !dbg [[DBG10]]
+; CHECK-NEXT:    br i1 [[OR_COND]], label %[[F]], label %[[E]], !dbg [[DBG14:![0-9]+]]
+; CHECK:       [[E]]:
+; CHECK-NEXT:    [[TMP0:%.*]] = tail call i32 (...) @g()
+; CHECK-NEXT:    br label %[[F]]
+; CHECK:       [[F]]:
+; CHECK-NEXT:    ret void
+;
+a:
+  br i1 %c0, label %b, label %c
+
+b:
+  br i1 %c1, label %f, label %d, !dbg !18
+
+c:
+  br i1 %c2, label %f, label %d, !dbg !18
+
+d:
+  %and = and i32 %x, %y, !dbg !19
+  %cmp = icmp eq i32 %and, 0, !dbg !20
+  br i1 %cmp, label %f, label %e, !dbg !21
+
+e:
+  %7 = tail call i32 (...) @g()
+  br label %f
+
+f:
+  ret void
+}
+
+!llvm.dbg.cu = !{!0}
+!llvm.debugify = !{!3, !4}
+!llvm.module.flags = !{!5}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C, file: !1, producer: "debugify", isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug, enums: !2)
+!1 = !DIFile(filename: "a.ll", directory: "/")
+!2 = !{}
+!3 = !{i32 9}
+!4 = !{i32 0}
+!5 = !{i32 2, !"Debug Info Version", i32 3}
+!7 = !DISubroutineType(types: !2)
+!17 = distinct !DISubprogram(name: "f", scope: null, file: !1, line: 1, type: !7, scopeLine: 1, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0, retainedNodes: !2)
+!18 = !DILocation(line: 10, column: 10, scope: !17)
+!19 = !DILocation(line: 10, column: 10, scope: !17, atomGroup: 1, atomRank: 2)
+!20 = !DILocation(line: 10, column: 10, scope: !17, atomGroup: 2, atomRank: 2)
+!21 = !DILocation(line: 10, column: 10, scope: !17, atomGroup: 2, atomRank: 1)
+;.
+; CHECK: [[META0:![0-9]+]] = distinct !DICompileUnit(language: DW_LANG_C, file: [[META1:![0-9]+]], producer: "debugify", isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug, enums: [[META2:![0-9]+]])
+; CHECK: [[META1]] = !DIFile(filename: "a.ll", directory: {{.*}})
+; CHECK: [[META2]] = !{}
+; CHECK: [[DBG6]] = distinct !DISubprogram(name: "f", scope: null, file: [[META1]], line: 1, type: [[META7:![0-9]+]], scopeLine: 1, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: [[META0]], retainedNodes: [[META2]])
+; CHECK: [[META7]] = !DISubroutineType(types: [[META2]])
+; CHECK: [[DBG8]] = !DILocation(line: 10, column: 10, scope: [[DBG6]], atomGroup: 5, atomRank: 2)
+; CHECK: [[DBG9]] = !DILocation(line: 10, column: 10, scope: [[DBG6]], atomGroup: 6, atomRank: 2)
+; CHECK: [[DBG10]] = !DILocation(line: 10, column: 10, scope: [[DBG6]])
+; CHECK: [[DBG11]] = !DILocation(line: 10, column: 10, scope: [[DBG6]], atomGroup: 6, atomRank: 1)
+; CHECK: [[DBG12]] = !DILocation(line: 10, column: 10, scope: [[DBG6]], atomGroup: 3, atomRank: 2)
+; CHECK: [[DBG13]] = !DILocation(line: 10, column: 10, scope: [[DBG6]], atomGroup: 4, atomRank: 2)
+; CHECK: [[DBG14]] = !DILocation(line: 10, column: 10, scope: [[DBG6]], atomGroup: 4, atomRank: 1)
+;.