[llvm] [mlir] [OMPIRBuilder] Don't leak a foreign debug loc into the taskwait call (PR #222586)

[Abid Qadeer via llvm-commits]

https://github.com/abidh created https://github.com/llvm/llvm-project/pull/222586

A kmp_depend_info array is hoisted to the entry block of the enclosing function. Pointing the builder into that block also adopts the location of what is already there, which belongs to whichever construct put it there rather than to the construct being emitted. restoreIP does not put the location back either, since it adopts the location of the instruction it lands on, so in createTaskwait the leak outlives the excursion and reaches the __kmpc_omp_taskwait_deps_51 call.

Use InsertPointGuard, which restores the location along with the insertion point.

Fixes https://github.com/llvm/llvm-project/issues/222044

>From ed9f5f6df238c30fbaca9a551393468b5ddacad0 Mon Sep 17 00:00:00 2001
From: Abid Qadeer <haqadeer at amd.com>
Date: Tue, 18 Aug 2026 19:36:03 +0100
Subject: [PATCH] [OMPIRBuilder] Don't leak a foreign debug location into the
 taskwait call

A kmp_depend_info array is hoisted to the entry block of the enclosing
function. Pointing the builder into that block also adopts the location of
what is already there, which belongs to whichever construct put it there
rather than to the construct being emitted. restoreIP does not put the
location back either, since it adopts the location of the instruction it
lands on, so in createTaskwait the leak outlives the excursion and reaches
the __kmpc_omp_taskwait_deps_51 call.

Use InsertPointGuard, which restores the location along with the insertion
point.

Fixes https://github.com/llvm/llvm-project/issues/222044

Co-authored-by: Cursor <cursoragent at cursor.com>
---
 llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp     | 29 ++++++++-------
 .../Target/LLVMIR/openmp-task-depend-loc.mlir | 37 +++++++++++++++++++
 2 files changed, 52 insertions(+), 14 deletions(-)
 create mode 100644 mlir/test/Target/LLVMIR/openmp-task-depend-loc.mlir

diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 3f80c858d033e..4373b7ce3fa8b 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -2397,14 +2397,15 @@ static Value *emitTaskDependencies(
   Type *DependInfo = OMPBuilder.DependInfo;
 
   Value *DepArray = nullptr;
-  OpenMPIRBuilder::InsertPointTy OldIP = Builder.saveIP();
-  Builder.SetInsertPoint(
-      OldIP.getBlock()->getParent()->getEntryBlock().getTerminator());
-
   Type *DepArrayTy = ArrayType::get(DependInfo, Dependencies.size());
-  DepArray = Builder.CreateAlloca(DepArrayTy, nullptr, ".dep.arr.addr");
-
-  Builder.restoreIP(OldIP);
+  {
+    // Use a InsertPointGuard to restore the location back along with the
+    // insertion point.
+    IRBuilderBase::InsertPointGuard IPGuard(Builder);
+    Builder.SetInsertPoint(
+        Builder.GetInsertBlock()->getParent()->getEntryBlock().getTerminator());
+    DepArray = Builder.CreateAlloca(DepArrayTy, nullptr, ".dep.arr.addr");
+  }
 
   for (const auto &[DepIdx, Dep] : enumerate(Dependencies)) {
     Value *Base =
@@ -2439,16 +2440,16 @@ void OpenMPIRBuilder::createTaskwait(const LocationDescription &Loc,
     DepArray = Dependencies.DepArray;
     NumDeps = Dependencies.NumDeps;
   } else if (!Dependencies.Deps.empty()) {
-    InsertPointTy OldIP = Builder.saveIP();
-    BasicBlock &entryBB =
-        Builder.GetInsertBlock()->getParent()->getEntryBlock();
-    Builder.SetInsertPoint(&entryBB, entryBB.getFirstInsertionPt());
-
     DepArrayTy = ArrayType::get(DependInfo, Dependencies.Deps.size());
-    DepArray = Builder.CreateAlloca(DepArrayTy, nullptr, ".dep.arr.addr");
     NumDeps = Builder.getInt32(Dependencies.Deps.size());
+    {
+      IRBuilderBase::InsertPointGuard IPGuard(Builder);
+      BasicBlock &entryBB =
+          Builder.GetInsertBlock()->getParent()->getEntryBlock();
+      Builder.SetInsertPoint(&entryBB, entryBB.getFirstInsertionPt());
+      DepArray = Builder.CreateAlloca(DepArrayTy, nullptr, ".dep.arr.addr");
+    }
 
-    Builder.restoreIP(OldIP);
     for (const auto &[DepIdx, Dep] : enumerate(Dependencies.Deps)) {
       Value *Base =
           Builder.CreateConstInBoundsGEP2_64(DepArrayTy, DepArray, 0, DepIdx);
diff --git a/mlir/test/Target/LLVMIR/openmp-task-depend-loc.mlir b/mlir/test/Target/LLVMIR/openmp-task-depend-loc.mlir
new file mode 100644
index 0000000000000..2a6c5801a233a
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-task-depend-loc.mlir
@@ -0,0 +1,37 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+llvm.func @task_then_taskwait(%x: !llvm.ptr, %y: !llvm.ptr) {
+  omp.task depend(taskdependout -> %x : !llvm.ptr) {
+    omp.terminator
+  } loc(#loc_task)
+  omp.taskwait depend(taskdependin -> %y : !llvm.ptr) loc(#loc_wait)
+  llvm.return
+} loc(#loc_fn)
+
+// The task is on line 7 and the taskwait on line 11. Each runtime call must
+// carry its own line.
+
+// CHECK: define void @task_then_taskwait
+// CHECK: call i32 @__kmpc_omp_task_with_deps({{.*}}), !dbg ![[TASK:[0-9]+]]
+// CHECK: call void @__kmpc_omp_taskwait_deps_51({{.*}}), !dbg ![[WAIT:[0-9]+]]
+// CHECK-DAG: ![[TASK]] = !DILocation(line: 7, column: 9
+// CHECK-DAG: ![[WAIT]] = !DILocation(line: 11, column: 9
+
+#di_file = #llvm.di_file<"test.f90" in "">
+#di_null_type = #llvm.di_null_type
+#di_compile_unit = #llvm.di_compile_unit<id = distinct[0]<>,
+  sourceLanguage = DW_LANG_Fortran95, file = #di_file, producer = "flang",
+  isOptimized = false, emissionKind = LineTablesOnly>
+#di_subroutine_type = #llvm.di_subroutine_type<
+  callingConvention = DW_CC_normal, types = #di_null_type>
+#di_subprogram = #llvm.di_subprogram<id = distinct[1]<>,
+  compileUnit = #di_compile_unit, scope = #di_file, name = "sub1",
+  file = #di_file, subprogramFlags = "Definition", type = #di_subroutine_type>
+
+#loc1 = loc("test.f90":3:1)
+#loc2 = loc("test.f90":7:9)
+#loc3 = loc("test.f90":11:9)
+
+#loc_fn = loc(fused<#di_subprogram>[#loc1])
+#loc_task = loc(fused<#di_subprogram>[#loc2])
+#loc_wait = loc(fused<#di_subprogram>[#loc3])