[OpenMP] Improve debug info generation for outlined function.

In OMPIRBuilder, a new function is created for the TargetOp. We also
create a new DISubprogram for it. All the variables that were in the
target region now have to be updated to have the correct scope. This
after the fact updating of debug information becomes very difficult in
certain cases.

This PR is making the change that OMPIRBuilder will not generate a
DISubprogram. The responsibility has been shifted to the frontend. This
allows us to simplify the updating of debug records.

The PR is made a bit more complicated by the the fact that in new
scheme, the debug location already points to the new DISubprogram by
the time it reaches convertOmpTarget. But we need some code generation
in the parent function so we have to carefully manage the debug
locations.

This builds on `llvm#138039` and fixes issue `llvm#134991`.

Author: abidh

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -6887,23 +6887,19 @@ static void FixupDebugInfoForOutlinedFunction(
   if (!NewSP)
     return;
 
-  DenseMap<const MDNode *, MDNode *> Cache;
   SmallDenseMap<DILocalVariable *, DILocalVariable *> RemappedVariables;
 
   auto GetUpdatedDIVariable = [&](DILocalVariable *OldVar, unsigned arg) {
-    auto NewSP = Func->getSubprogram();
     DILocalVariable *&NewVar = RemappedVariables[OldVar];
     // Only use cached variable if the arg number matches. This is important
     // so that DIVariable created for privatized variables are not discarded.
     if (NewVar && (arg == NewVar->getArg()))
       return NewVar;
 
-    DILocalScope *NewScope = DILocalScope::cloneScopeForSubprogram(
-        *OldVar->getScope(), *NewSP, Builder.getContext(), Cache);
     NewVar = llvm::DILocalVariable::get(
-        Builder.getContext(), NewScope, OldVar->getName(), OldVar->getFile(),
-        OldVar->getLine(), OldVar->getType(), arg, OldVar->getFlags(),
-        OldVar->getAlignInBits(), OldVar->getAnnotations());
+        Builder.getContext(), OldVar->getScope(), OldVar->getName(),
+        OldVar->getFile(), OldVar->getLine(), OldVar->getType(), arg,
+        OldVar->getFlags(), OldVar->getAlignInBits(), OldVar->getAnnotations());
     return NewVar;
   };
 
@@ -6917,7 +6913,8 @@ static void FixupDebugInfoForOutlinedFunction(
         ArgNo = std::get<1>(Iter->second) + 1;
       }
     }
-    DR->setVariable(GetUpdatedDIVariable(OldVar, ArgNo));
+    if (ArgNo != 0)
+      DR->setVariable(GetUpdatedDIVariable(OldVar, ArgNo));
   };
 
   // The location and scope of variable intrinsics and records still point to
@@ -6996,36 +6993,9 @@ static Expected<Function *> createOutlinedFunction(
 
   // Save insert point.
   IRBuilder<>::InsertPointGuard IPG(Builder);
-  // If there's a DISubprogram associated with current function, then
-  // generate one for the outlined function.
-  if (Function *ParentFunc = BB->getParent()) {
-    if (DISubprogram *SP = ParentFunc->getSubprogram()) {
-      DICompileUnit *CU = SP->getUnit();
-      DIBuilder DB(*M, true, CU);
-      DebugLoc DL = Builder.getCurrentDebugLocation();
-      if (DL) {
-        // TODO: We are using nullopt for arguments at the moment. This will
-        // need to be updated when debug data is being generated for variables.
-        DISubroutineType *Ty =
-            DB.createSubroutineType(DB.getOrCreateTypeArray({}));
-        DISubprogram::DISPFlags SPFlags = DISubprogram::SPFlagDefinition |
-                                          DISubprogram::SPFlagOptimized |
-                                          DISubprogram::SPFlagLocalToUnit;
-
-        DISubprogram *OutlinedSP = DB.createFunction(
-            CU, FuncName, FuncName, SP->getFile(), DL.getLine(), Ty,
-            DL.getLine(), DINode::DIFlags::FlagArtificial, SPFlags);
-
-        // Attach subprogram to the function.
-        Func->setSubprogram(OutlinedSP);
-        // Update the CurrentDebugLocation in the builder so that right scope
-        // is used for things inside outlined function.
-        Builder.SetCurrentDebugLocation(
-            DILocation::get(Func->getContext(), DL.getLine(), DL.getCol(),
-                            OutlinedSP, DL.getInlinedAt()));
-      }
-    }
-  }
+  // We will generate the entries in the outlined function but the debug
+  // location may still be pointing to the parent function. Reset it now.
+  Builder.SetCurrentDebugLocation(llvm::DebugLoc());
 
   // Generate the region into the function.
   BasicBlock *EntryBB = BasicBlock::Create(Builder.getContext(), "entry", Func);

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

@@ -5299,6 +5299,20 @@ static LogicalResult
 convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
                  LLVM::ModuleTranslation &moduleTranslation) {
   auto targetOp = cast<omp::TargetOp>(opInst);
+  // The current debug location already has the DISubprogram for the outlined
+  // function that will be created for the target op. We save it here so that
+  // we can set it on the outlined function.
+  llvm::DebugLoc OutlinedFnLoc = builder.getCurrentDebugLocation();
+  // During the handling of target op, we will generate instructions in the
+  // parent function like call to oulined function or branch to new BasicBlock.
+  // We set the debug location here to parent function so that those get the
+  // correct debug locations. For outlined functions, the normal MLIR op
+  // conversion will automatically pick the correct location.
+  llvm::BasicBlock *parentBB = builder.GetInsertBlock();
+  if (parentBB && !parentBB->empty())
+    builder.SetCurrentDebugLocation(parentBB->back().getDebugLoc());
+  else
+    builder.SetCurrentDebugLocation(llvm::DebugLoc());
   if (failed(checkImplementationStatus(opInst)))
     return failure();
 
@@ -5395,6 +5409,9 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
     assert(llvmParentFn && llvmOutlinedFn &&
            "Both parent and outlined functions must exist at this point");
 
+    if (OutlinedFnLoc && llvmParentFn->getSubprogram())
+      llvmOutlinedFn->setSubprogram(OutlinedFnLoc->getScope()->getSubprogram());
+
     if (auto attr = llvmParentFn->getFnAttribute("target-cpu");
         attr.isStringAttribute())
       llvmOutlinedFn->addFnAttr(attr);

mlir/test/Target/LLVMIR/omptarget-debug-nowait.mlir

@@ -20,7 +20,7 @@ module attributes {omp.is_target_device = false} {
     ^bb3:  // pred: ^bb1
       llvm.store %13, %arg1 : i32, !llvm.ptr
       omp.terminator
-    }
+    } loc(#loc3)
     llvm.return
   } loc(#loc2)
 }
@@ -34,7 +34,10 @@ module attributes {omp.is_target_device = false} {
  types = #di_null_type>
 #sp = #llvm.di_subprogram<compileUnit = #cu, name = "main", file=#file,
  subprogramFlags = "Definition", type = #sp_ty>
+#sp1 = #llvm.di_subprogram<compileUnit = #cu, name = "target", file=#file,
+ subprogramFlags = "Definition", type = #sp_ty>
 
 #loc1 = loc("test.f90":6:7)
 #loc2 = loc(fused<#sp>[#loc1])
+#loc3 = loc(fused<#sp1>[#loc1])
 

mlir/test/Target/LLVMIR/omptarget-debug-var-1.mlir

@@ -19,11 +19,13 @@
 #g_var_expr = #llvm.di_global_variable_expression<var = #g_var>
 #sp = #llvm.di_subprogram<id = distinct[2]<>, compileUnit = #cu, scope = #file,
   name = "test", file = #file, subprogramFlags = "Definition", type = #sp_ty>
-#var_arr = #llvm.di_local_variable<scope = #sp,
+#sp1 = #llvm.di_subprogram<id = distinct[3]<>, compileUnit = #cu, scope = #file,
+  name = "target", file = #file, subprogramFlags = "Definition", type = #sp_ty>
+#var_arr = #llvm.di_local_variable<scope = #sp1,
   name = "arr", file = #file, line = 4, type = #array_ty>
-#var_i = #llvm.di_local_variable<scope = #sp,
+#var_i = #llvm.di_local_variable<scope = #sp1,
   name = "i", file = #file, line = 13, type = #int_ty>
-#var_x = #llvm.di_local_variable<scope = #sp,
+#var_x = #llvm.di_local_variable<scope = #sp1,
  name = "x", file = #file, line = 12, type = #real_ty>
 
 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} {
@@ -47,7 +49,7 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memo
       llvm.intr.dbg.declare #var_arr = %arg1 : !llvm.ptr
       llvm.intr.dbg.declare #var_i = %arg2 : !llvm.ptr
       omp.terminator
-    }
+    } loc(#loc5)
     llvm.return
   } loc(#loc3)
   llvm.mlir.global internal @_QFEarr() {addr_space = 0 : i32, dbg_exprs = [#g_var_expr]} : !llvm.array<10 x i32> {
@@ -57,8 +59,9 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memo
 #loc2 = loc("target.f90":11:7)
 #loc3 = loc(fused<#sp>[#loc2])
 #loc4 = loc(fused<#g_var>[#loc1])
+#loc5 = loc(fused<#sp1>[#loc2])
 
-// CHECK: ![[SP:[0-9]+]] = distinct !DISubprogram(name: "__omp_offloading{{.*}}test{{.*}})
+// CHECK: ![[SP:[0-9]+]] = distinct !DISubprogram(name: "target"{{.*}})
 // CHECK: !DILocalVariable(name: "dyn_ptr", arg: 1, scope: ![[SP]]{{.*}}flags: DIFlagArtificial)
 // CHECK: !DILocalVariable(name: "x", arg: 2, scope: ![[SP]]{{.*}})
 // CHECK: !DILocalVariable(name: "arr", arg: 3, scope: ![[SP]]{{.*}})

mlir/test/Target/LLVMIR/omptarget-debug-var-2.mlir

@@ -19,11 +19,13 @@
 #g_var_expr = #llvm.di_global_variable_expression<var = #g_var>
 #sp = #llvm.di_subprogram<id = distinct[2]<>, compileUnit = #cu, scope = #file,
   name = "test", file = #file, subprogramFlags = "Definition", type = #sp_ty>
-#var_arr = #llvm.di_local_variable<scope = #sp,
+#sp1 = #llvm.di_subprogram<id = distinct[3]<>, compileUnit = #cu, scope = #file,
+  name = "target", file = #file, subprogramFlags = "Definition", type = #sp_ty>
+#var_arr = #llvm.di_local_variable<scope = #sp1,
   name = "arr", file = #file, line = 4, type = #array_ty>
-#var_i = #llvm.di_local_variable<scope = #sp,
+#var_i = #llvm.di_local_variable<scope = #sp1,
   name = "i", file = #file, line = 13, type = #int_ty>
-#var_x = #llvm.di_local_variable<scope = #sp,
+#var_x = #llvm.di_local_variable<scope = #sp1,
  name = "x", file = #file, line = 12, type = #real_ty>
 
 module attributes {omp.is_target_device = false} {
@@ -45,7 +47,7 @@ module attributes {omp.is_target_device = false} {
       llvm.intr.dbg.declare #var_arr = %arg1 : !llvm.ptr
       llvm.intr.dbg.declare #var_i = %arg2 : !llvm.ptr
       omp.terminator
-    }
+    }  loc(#loc5)
     llvm.return
   } loc(#loc3)
   llvm.mlir.global internal @_QFEarr() {addr_space = 0 : i32, dbg_exprs = [#g_var_expr]} : !llvm.array<10 x i32> {
@@ -55,8 +57,9 @@ module attributes {omp.is_target_device = false} {
 #loc2 = loc("target.f90":11:7)
 #loc3 = loc(fused<#sp>[#loc2])
 #loc4 = loc(fused<#g_var>[#loc1])
+#loc5 = loc(fused<#sp1>[#loc2])
 
-// CHECK: ![[SP:[0-9]+]] = distinct !DISubprogram(name: "__omp_offloading{{.*}}test{{.*}})
+// CHECK: ![[SP:[0-9]+]] = distinct !DISubprogram(name: "target"{{.*}})
 // CHECK: !DILocalVariable(name: "x", arg: 1, scope: ![[SP]]{{.*}})
 // CHECK: !DILocalVariable(name: "arr", arg: 2, scope: ![[SP]]{{.*}})
 // CHECK: !DILocalVariable(name: "i", arg: 3, scope: ![[SP]]{{.*}})

mlir/test/Target/LLVMIR/omptarget-debug.mlir

@@ -10,7 +10,7 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memo
       %13 = llvm.mlir.constant(1 : i32) : i32
       llvm.store %13, %arg0 : i32, !llvm.ptr loc(#loc2)
       omp.terminator
-    }
+    } loc(#loc4)
     llvm.return
   } loc(#loc3)
 }
@@ -21,9 +21,13 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memo
 #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":1:1)
 #loc2 = loc("target.f90":46:3)
 #loc3 = loc(fused<#sp>[#loc1])
+#loc4 = loc(fused<#sp1>[#loc1])
 
-// CHECK-DAG: ![[SP:.*]] = {{.*}}!DISubprogram(name: "__omp_offloading_{{.*}}"{{.*}})
+// CHECK-DAG: ![[SP:.*]] = {{.*}}!DISubprogram(name: "__omp_offloading_target"{{.*}})
 // CHECK-DAG: !DILocation(line: 46, column: 3, scope: ![[SP]])

mlir/test/Target/LLVMIR/omptarget-debug2.mlir

@@ -11,7 +11,7 @@ module attributes {omp.is_target_device = false} {
       %13 = llvm.mlir.constant(1 : i32) : i32
       llvm.store %13, %arg0 : i32, !llvm.ptr loc(#loc2)
       omp.terminator
-    }
+    }  loc(#loc4)
     llvm.return
   } loc(#loc3)
 }
@@ -22,9 +22,13 @@ module attributes {omp.is_target_device = false} {
 #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":1:1)
 #loc2 = loc("target.f90":46:3)
 #loc3 = loc(fused<#sp>[#loc1])
+#loc4 = loc(fused<#sp1>[#loc1])
 
-// CHECK-DAG: ![[SP:.*]] = {{.*}}!DISubprogram(name: "__omp_offloading_{{.*}}"{{.*}})
+// CHECK-DAG: ![[SP:.*]] = {{.*}}!DISubprogram(name: "__omp_offloading_target"{{.*}})
 // CHECK-DAG: !DILocation(line: 46, column: 3, scope: ![[SP]])

mlir/test/Target/LLVMIR/omptarget-parallel-llvm-debug.mlir

@@ -6,8 +6,10 @@
 #cu = #llvm.di_compile_unit<id = distinct[0]<>, sourceLanguage = DW_LANG_Fortran95, file = #di_file, producer = "flang", isOptimized = false, emissionKind = Full>
 #sp_ty = #llvm.di_subroutine_type<callingConvention = DW_CC_program, types = #di_null_type>
 #sp = #llvm.di_subprogram<compileUnit = #cu, scope = #di_file, name = "test", file = #di_file, subprogramFlags = "Definition", type = #sp_ty>
+#sp1 = #llvm.di_subprogram<compileUnit = #cu, scope = #di_file, name = "kernel", file = #di_file, subprogramFlags = "Definition", type = #sp_ty>
 #int_ty = #llvm.di_basic_type<tag = DW_TAG_base_type, name = "integer", sizeInBits = 32, encoding = DW_ATE_signed>
 #var_x = #llvm.di_local_variable<scope = #sp, name = "x", file = #di_file, type = #int_ty>
+#var_x1 = #llvm.di_local_variable<scope = #sp1, name = "x", file = #di_file, type = #int_ty>
 module attributes {dlti.dl_spec = #dlti.dl_spec<i32 = dense<32> : vector<2xi64>, f16 = dense<16> : vector<2xi64>, f64 = dense<64> : vector<2xi64>, f128 = dense<128> : vector<2xi64>, i128 = dense<128> : vector<2xi64>, i64 = dense<64> : vector<2xi64>, f80 = dense<128> : vector<2xi64>, !llvm.ptr<272> = dense<64> : vector<4xi64>, !llvm.ptr<271> = dense<32> : vector<4xi64>, !llvm.ptr = dense<64> : vector<4xi64>, !llvm.ptr<270> = dense<32> : vector<4xi64>, i1 = dense<8> : vector<2xi64>, i8 = dense<8> : vector<2xi64>, i16 = dense<16> : vector<2xi64>, "dlti.endianness" = "little", "dlti.stack_alignment" = 128 : i64, "dlti.mangling_mode" = "e">, fir.defaultkind = "a1c4d8i4l4r4", fir.kindmap = "", fir.target_cpu = "x86-64", llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128", llvm.ident = "flang version 21.0.0 (/home/haqadeer/work/src/aomp-llvm-project/flang 793f9220ab32f92fc3b253efec2e332c18090e53)", llvm.target_triple = "x86_64-unknown-linux-gnu", omp.is_gpu = false, omp.is_target_device = false, omp.requires = #omp<clause_requires none>, omp.target_triples = ["amdgcn-amd-amdhsa"], omp.version = #omp.version<version = 52>} {
   llvm.func @_QQmain() attributes {fir.bindc_name = "test", frame_pointer = #llvm.framePointerKind<all>, target_cpu = "x86-64"} {
     %0 = llvm.mlir.constant(1 : i64) : i64
@@ -16,20 +18,22 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<i32 = dense<32> : vector<2xi64>,
     %5 = omp.map.info var_ptr(%1 : !llvm.ptr, i32) map_clauses(implicit, exit_release_or_enter_alloc) capture(ByCopy) -> !llvm.ptr {name = "x"}
     omp.target map_entries(%5 -> %arg0 : !llvm.ptr) {
       %6 = llvm.mlir.constant(1 : i32) : i32
-      llvm.intr.dbg.declare #var_x = %arg0 : !llvm.ptr loc(#loc2)
+      llvm.intr.dbg.declare #var_x1 = %arg0 : !llvm.ptr loc(#loc3)
       omp.parallel {
         %7 = llvm.load %arg0 : !llvm.ptr -> i32
         %8 = llvm.add %7, %6 : i32
         llvm.store %8, %arg0 : i32, !llvm.ptr
         omp.terminator
       }
       omp.terminator
-    }
+    } loc(#loc4)
     llvm.return
   } loc(#loc10)
 }
 #loc1 = loc("target.f90":1:7)
 #loc2 = loc("target.f90":3:18)
+#loc3 = loc("target.f90":6:18)
+#loc4 = loc(fused<#sp1>[#loc3])
 #loc10 = loc(fused<#sp>[#loc1])