[Mlir-commits] [llvm] [mlir] [OMPIRBuilder] Keep the target construct's location for the deinit call (PR #221303)

[Abid Qadeer]

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

#217407 gave __kmpc_target_init and __kmpc_target_deinit a debug location by setting the outlined function's location once, before generating the target body. The init call is emitted before the body and keeps it, but the deinit call is emitted after, and by then the body may have replaced it.

The body callback builds the body with its own IRBuilder, so nothing it emits can disturb the one createOutlinedFunction() uses for the epilogue. But a body holding another OpenMP construct, a nested parallel say, calls OpenMPIRBuilder::createParallel, and that leaves the builder pointing at the wrong debug location, or at none at all. The deinit call then silently loses its !dbg again.

The test added with #217407 uses a target region whose body is a single store. Such a body never calls back into OpenMPIRBuilder, the location survives, and the gap was not visible.

Re-establish the outlined function's location after the body callback, before emitting the epilogue. The epilogue belongs to the target construct rather than to whatever the body emitted last, so this is the location it should carry, and init and deinit go back to agreeing.

Extends the existing test with a nested parallel, which fails on the deinit check without this change.

>From 129a5b993e321138f50524175b9ef4f7b3fcf02a Mon Sep 17 00:00:00 2001
From: Abid Qadeer <haqadeer at amd.com>
Date: Fri, 4 Sep 2026 18:30:38 +0100
Subject: [PATCH] [OMPIRBuilder] Keep the target construct's location for the
 deinit call

#217407 gave __kmpc_target_init and __kmpc_target_deinit a debug location by
setting the outlined function's location once, before generating the target
body. The init call is emitted before the body and keeps it, but the deinit
call is emitted after, and by then the body may have replaced it.

The body callback builds the body with its own IRBuilder, so nothing it emits
can disturb the one createOutlinedFunction() uses for the epilogue. But a body
holding another OpenMP construct, a nested parallel say, calls
OpenMPIRBuilder::createParallel, and that leaves the builder pointing at the
wrong debug location, or at none at all. The deinit call then silently loses
its !dbg again.

The test added with #217407 uses a target region whose body is a single store.
Such a body never calls back into OpenMPIRBuilder, the location survives, and
the gap was not visible.

Re-establish the outlined function's location after the body callback, before
emitting the epilogue. The epilogue belongs to the target construct rather
than to whatever the body emitted last, so this is the location it should
carry, and init and deinit go back to agreeing.

Extends the existing test with a nested parallel, which fails on the deinit
check without this change.

Co-authored-by: Cursor <cursoragent at cursor.com>
---
 llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp     |  8 ++++
 .../omptarget-debug-runtime-call-loc.mlir     | 46 ++++++++++++++++++-
 2 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index a168b8651d1fe..8f88f4ba59319 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -9306,6 +9306,14 @@ static Expected<Function *> createOutlinedFunction(
   if (!AfterIP)
     return AfterIP.takeError();
   Builder.SetInsertPoint(ExitBB);
+  // The body callback builds the body with its own IRBuilder and cannot reach
+  // this one directly. But a body holding another OpenMP construct, a nested
+  // parallel say, calls OpenMPIRBuilder::createParallel, and that can leave
+  // this Builder pointing at the wrong debug location, or at none at all. The
+  // epilogue below belongs to the target construct rather than to whatever the
+  // body emitted last, so re-establish the location the prologue was emitted
+  // with.
+  Builder.SetCurrentDebugLocation(OutlinedFnLoc);
 
   // Insert target deinit call in the device compilation pass.
   if (OMPBuilder.Config.isTargetDevice())
diff --git a/mlir/test/Target/LLVMIR/omptarget-debug-runtime-call-loc.mlir b/mlir/test/Target/LLVMIR/omptarget-debug-runtime-call-loc.mlir
index 3baed4c9b36fb..d78fd24488da4 100644
--- a/mlir/test/Target/LLVMIR/omptarget-debug-runtime-call-loc.mlir
+++ b/mlir/test/Target/LLVMIR/omptarget-debug-runtime-call-loc.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck %s
 
 // The __kmpc_target_init and __kmpc_target_deinit calls are emitted into the
 // outlined kernel by OpenMPIRBuilder, which has no location of its own for
@@ -41,3 +41,47 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memo
 // CHECK: call void @__kmpc_target_deinit(), !dbg ![[LOC]]
 // CHECK-DAG: ![[SP:[0-9]+]] = distinct !DISubprogram(name: "__omp_offloading_target"
 // CHECK-DAG: ![[LOC]] = !DILocation(line: 12, column: 5, scope: ![[SP]])
+
+// -----
+
+// The deinit call is emitted after the target body has been generated, so it
+// picks up whatever debug location the body left behind. A body holding
+// another construct leaves it empty, which used to cost the deinit its !dbg
+// while the init, emitted before the body, kept its own.
+
+module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui32>>, llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_target_device = true} {
+  llvm.func @_QQmain() {
+    %0 = llvm.mlir.constant(1 : i32) : i32
+    %1 = llvm.alloca %0 x i32 : (i32) -> !llvm.ptr<5>
+    %ascast = llvm.addrspacecast %1 : !llvm.ptr<5> to !llvm.ptr
+    %9 = omp.map.info var_ptr(%ascast : !llvm.ptr, i32) map_clauses(tofrom) capture(ByRef) name("") -> !llvm.ptr
+    omp.target kernel_type(generic) map_entries(%9 -> %arg0 : !llvm.ptr) {
+      omp.parallel {
+        %13 = llvm.mlir.constant(1 : i32) : i32
+        llvm.store %13, %arg0 : i32, !llvm.ptr loc(#loc2)
+        omp.terminator
+      } loc(#loc4)
+      omp.terminator
+    } loc(#loc4)
+    llvm.return
+  } loc(#loc3)
+}
+#file = #llvm.di_file<"target.f90" in "">
+#cu = #llvm.di_compile_unit<id = distinct[0]<>,
+ sourceLanguage = DW_LANG_Fortran95, file = #file, isOptimized = false,
+ emissionKind = LineTablesOnly>
+#sp_ty = #llvm.di_subroutine_type<callingConvention = DW_CC_normal>
+#sp = #llvm.di_subprogram<id = distinct[1]<>, compileUnit = #cu, scope = #file,
+ name = "_QQmain", file = #file, subprogramFlags = "Definition", type = #sp_ty>
+#sp1 = #llvm.di_subprogram<id = distinct[2]<>, compileUnit = #cu, scope = #file,
+ name = "__omp_offloading_target", file = #file, subprogramFlags = "Definition",
+ type = #sp_ty>
+#loc1 = loc("target.f90":12:5)
+#loc2 = loc("target.f90":46:3)
+#loc3 = loc(fused<#sp>[#loc1])
+#loc4 = loc(fused<#sp1>[#loc1])
+
+// CHECK: call i32 @__kmpc_target_init({{.*}}), !dbg ![[PLOC:[0-9]+]]
+// CHECK: call void @__kmpc_target_deinit(), !dbg ![[PLOC]]
+// CHECK-DAG: ![[PSP:[0-9]+]] = distinct !DISubprogram(name: "__omp_offloading_target"
+// CHECK-DAG: ![[PLOC]] = !DILocation(line: 12, column: 5, scope: ![[PSP]])