From b45ad3804a3f4b81f1434be1887064824e11d1bc Mon Sep 17 00:00:00 2001
From: Orlando Cazalet-Hyams <orlando.hyams@sony.com>
Date: Thu, 8 May 2025 13:29:46 +0100
Subject: [PATCH] [KeyInstr] Don't propagate source atoms to new uncond br in
 splitBasicBlock

splitBasicBlock inserts an unconditional branch in the "before" block to the
"after" block. It copies the DebugLoc from the split point. Prevent it copying
the source location atom.

Add unittest.

RFC: https://discourse.llvm.org/t/rfc-improving-is-stmt-placement-for-better-interactive-debugging/82668
---
 llvm/lib/IR/BasicBlock.cpp                  |  4 +-
 llvm/unittests/IR/BasicBlockDbgInfoTest.cpp | 69 +++++++++++++++++++++
 2 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/IR/BasicBlock.cpp b/llvm/lib/IR/BasicBlock.cpp
index b4c16447bc686..aea07f4731e7d 100644
--- a/llvm/lib/IR/BasicBlock.cpp
+++ b/llvm/lib/IR/BasicBlock.cpp
@@ -612,7 +612,7 @@ BasicBlock *BasicBlock::splitBasicBlock(iterator I, const Twine &BBName,
                                        this->getNextNode());
 
   // Save DebugLoc of split point before invalidating iterator.
-  DebugLoc Loc = I->getStableDebugLoc();
+  DebugLoc Loc = I->getStableDebugLoc()->getWithoutAtom();
   // Move all of the specified instructions from the original basic block into
   // the new basic block.
   New->splice(New->end(), this, I, end());
@@ -641,7 +641,7 @@ BasicBlock *BasicBlock::splitBasicBlockBefore(iterator I, const Twine &BBName) {
 
   BasicBlock *New = BasicBlock::Create(getContext(), BBName, getParent(), this);
   // Save DebugLoc of split point before invalidating iterator.
-  DebugLoc Loc = I->getDebugLoc();
+  DebugLoc Loc = I->getDebugLoc()->getWithoutAtom();
   // Move all of the specified instructions from the original basic block into
   // the new basic block.
   New->splice(New->end(), this, begin(), I);
diff --git a/llvm/unittests/IR/BasicBlockDbgInfoTest.cpp b/llvm/unittests/IR/BasicBlockDbgInfoTest.cpp
index 5607c633b7a88..cb1a56a26f063 100644
--- a/llvm/unittests/IR/BasicBlockDbgInfoTest.cpp
+++ b/llvm/unittests/IR/BasicBlockDbgInfoTest.cpp
@@ -149,6 +149,75 @@ TEST(BasicBlockDbgInfoTest, SplitBasicBlockBefore) {
   ASSERT_TRUE(I2->hasDbgRecords());
 }
 
+TEST(BasicBlockDbgInfoTest, DropSourceAtomOnSplit) {
+  LLVMContext C;
+  std::unique_ptr<Module> M = parseIR(C, R"---(
+    define dso_local void @func() !dbg !10 {
+      %1 = alloca i32, align 4
+      ret void, !dbg !DILocation(line: 3, column: 2, scope: !10, atomGroup: 1, atomRank: 1)
+    }
+
+    !llvm.dbg.cu = !{!0}
+    !llvm.module.flags = !{!2, !3}
+
+    !0 = distinct !DICompileUnit(language: DW_LANG_C11, file: !1, producer: "dummy", isOptimized: false, runtimeVersion: 0, emissionKind: FullDebug, splitDebugInlining: false, nameTableKind: None)
+    !1 = !DIFile(filename: "dummy", directory: "dummy")
+    !2 = !{i32 7, !"Dwarf Version", i32 5}
+    !3 = !{i32 2, !"Debug Info Version", i32 3}
+    !10 = distinct !DISubprogram(name: "func", scope: !1, file: !1, line: 1, type: !11, scopeLine: 1, spFlags: DISPFlagDefinition, unit: !0, retainedNodes: !13)
+    !11 = !DISubroutineType(types: !12)
+    !12 = !{null}
+    !13 = !{}
+    !14 = !DILocalVariable(name: "a", scope: !10, file: !1, line: 2, type: !15)
+    !15 = !DIBasicType(name: "int", size: 32, encoding: DW_ATE_signed)
+  )---");
+  ASSERT_TRUE(M);
+
+  Function *F = M->getFunction("func");
+
+  // Test splitBasicBlockBefore.
+  {
+    BasicBlock &BB = F->back();
+    // Split at `ret void`.
+    BasicBlock *Before =
+        BB.splitBasicBlockBefore(std::prev(BB.end(), 1), "before");
+    const DebugLoc &BrToAfterDL = Before->getTerminator()->getDebugLoc();
+    ASSERT_TRUE(BrToAfterDL);
+    EXPECT_EQ(BrToAfterDL->getAtomGroup(), 0u);
+
+    BasicBlock *After = Before->getSingleSuccessor();
+    ASSERT_TRUE(After);
+    const DebugLoc &OrigTerminatorDL = After->getTerminator()->getDebugLoc();
+    ASSERT_TRUE(OrigTerminatorDL);
+#ifdef EXPERIMENTAL_KEY_INSTRUCTIONS
+    EXPECT_EQ(OrigTerminatorDL->getAtomGroup(), 1u);
+#else
+    EXPECT_EQ(OrigTerminatorDL->getAtomGroup(), 0u);
+#endif
+  }
+
+  // Test splitBasicBlock.
+  {
+    BasicBlock &BB = F->back();
+    // Split at `ret void`.
+    BasicBlock *After = BB.splitBasicBlock(std::prev(BB.end(), 1), "before");
+
+    const DebugLoc &OrigTerminatorDL = After->getTerminator()->getDebugLoc();
+    ASSERT_TRUE(OrigTerminatorDL);
+#ifdef EXPERIMENTAL_KEY_INSTRUCTIONS
+    EXPECT_EQ(OrigTerminatorDL->getAtomGroup(), 1u);
+#else
+    EXPECT_EQ(OrigTerminatorDL->getAtomGroup(), 0u);
+#endif
+
+    BasicBlock *Before = After->getSinglePredecessor();
+    ASSERT_TRUE(Before);
+    const DebugLoc &BrToAfterDL = Before->getTerminator()->getDebugLoc();
+    ASSERT_TRUE(BrToAfterDL);
+    EXPECT_EQ(BrToAfterDL->getAtomGroup(), 0u);
+  }
+}
+
 TEST(BasicBlockDbgInfoTest, MarkerOperations) {
   LLVMContext C;
   std::unique_ptr<Module> M = parseIR(C, R"(