[Mlir-commits] [llvm] [mlir] [OMPIRBuilder] Move debug records to correct blocks. (PR #157125)

[Scott Linder]

================
@@ -8363,29 +8350,54 @@ static void FixupDebugInfoForOutlinedFunction(
         OldVar->getFlags(), OldVar->getAlignInBits(), OldVar->getAnnotations());
     return NewVar;
   };
-
-  auto UpdateDebugRecord = [&](auto *DR) {
+  auto UpdateDebugRecord = [&](DbgVariableRecord *DR) {
     DILocalVariable *OldVar = DR->getVariable();
     unsigned ArgNo = 0;
-    for (auto Loc : DR->location_ops()) {
-      auto Iter = ValueReplacementMap.find(Loc);
-      if (Iter != ValueReplacementMap.end()) {
-        DR->replaceVariableLocationOp(Loc, std::get<0>(Iter->second));
-        ArgNo = std::get<1>(Iter->second) + 1;
-      }
+    if (DR->getNumVariableLocationOps() != 1u)
+      return;
+    auto Loc = DR->getVariableLocationOp(0u);
+    auto Iter = ValueReplacementMap.find(Loc);
+    if (Iter != ValueReplacementMap.end()) {
+      DR->replaceVariableLocationOp(Loc, std::get<0>(Iter->second));
+      ArgNo = std::get<1>(Iter->second) + 1;
     }
     if (ArgNo != 0)
       DR->setVariable(GetUpdatedDIVariable(OldVar, ArgNo));
   };
 
+  SmallVector<DbgVariableRecord *, 4> DVRsToDelete;
+  auto MoveDebugRecordToCorrectBlock = [&](DbgVariableRecord *DVR) {
+    if (DVR->getNumVariableLocationOps() != 1u)
+      return;
+    auto Loc = DVR->getVariableLocationOp(0u);
+    BasicBlock *CurBB = DVR->getParent();
+    BasicBlock *RequiredBB = nullptr;
+
+    if (Instruction *LocInst = dyn_cast<Instruction>(Loc))
+      RequiredBB = LocInst->getParent();
+    else if (isa<llvm::Argument>(Loc))
+      RequiredBB = &(DVR->getFunction()->getEntryBlock());
+
+    if (RequiredBB && RequiredBB != CurBB) {
+      assert(!RequiredBB->empty());
+      RequiredBB->insertDbgRecordBefore(DVR->clone(),
+                                        RequiredBB->back().getIterator());
+      DVRsToDelete.push_back(DVR);
+    }
+  };
+
   // The location and scope of variable intrinsics and records still point to
   // the parent function of the target region. Update them.
   for (Instruction &I : instructions(Func)) {
     assert(!isa<llvm::DbgVariableIntrinsic>(&I) &&
            "Unexpected debug intrinsic");
-    for (DbgVariableRecord &DVR : filterDbgVars(I.getDbgRecordRange()))
+    for (DbgVariableRecord &DVR : filterDbgVars(I.getDbgRecordRange())) {
       UpdateDebugRecord(&DVR);
+      MoveDebugRecordToCorrectBlock(&DVR);
----------------
slinder1 wrote:

You're right, I'm not sure what performance I was talking about!

The lambdas LGTM

https://github.com/llvm/llvm-project/pull/157125