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

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-openmp

Author: Abid Qadeer (abidh)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/222586.diff


2 Files Affected:

- (modified) llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp (+15-14) 
- (added) mlir/test/Target/LLVMIR/openmp-task-depend-loc.mlir (+37) 


``````````diff
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])

``````````

</details>


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