[OMPIRBuilder] Avoid invalid debug location. (#151306)

This fixes #147063.

I tried to fix this issue in more general way in
#147091 but the reviewer
suggested to fix the locations which are causing this issue. So this is
a more targeted approach.

The `restoreIP` is frequently used in the `OMPIRBuilder` to change the
insert position. This function eventually calls
`SetInsertPoint(BasicBlock *TheBB, BasicBlock::iterator IP)`. This
function updates the insert point and the debug location. But if the
`IP` is pointing to the end of the `TheBB`, then the debug location is
not updated and we could have a mismatch between insert point and the
debug location.

The problem can occur in 2 different code patterns. 

This code below shows the first scenario.
```
  1. auto curPos = builder.saveIP();
  2. builder.restoreIP(/* some new pos */);
  3.  // generate some code
  4. builder.restoreIP(curPos);
```

If `curPos` points to the end of basic block, we could have a problem.
But it is easy one to handle as we have the location before hand and can
save the correct debug location before 2 and then restore it after 3.
This can be done either manually or using the `llvm::InsertPointGuard`
as shown below.

```
  // manual approach
  auto curPos = builder.saveIP();
  llvm::DebugLoc DbgLoc = builder.getCurrentDebugLocation();
  builder.restoreIP(/* some new pos */);
  // generate some code
  builder.SetCurrentDebugLocation(DbgLoc);
  builder.restoreIP(curPos);

  {
     // using InsertPointGuard
     llvm::InsertPointGuard IPG(builder);
     builder.restoreIP(/* some new pos */);
     // generate some code
  }
```
This PR fixes one problematic case using the manual approach.

For the 2nd scenario, look at the code below.

```
1. void fn(InsertPointTy allocIP, InsertPointTy codegenIP) {
2.   builder.setInsertPoint(allocIP);
3.   // generate some alloca
4.   builder.setInsertPoint(codegenIP);
5. }
```

The `fn` can be called from anywhere and we can't assume the debug
location of the builder is valid at the start of the function. So if 4
does not update the debug location because the `codegenIP` points at the
end of the block, the rest of the code can end up using the debug
location of the `allocaIP`. Unlike the first case, we don't have a debug
location that we can save before hand and restore afterwards.

The solution here is to use the location of the last instruction in that
block. I have added a wrapper function over `restoreIP` that could be
called for such cases. This PR uses it to fix one problematic case.

Author: abidh

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -92,6 +92,18 @@ static bool isConflictIP(IRBuilder<>::InsertPoint IP1,
   return IP1.getBlock() == IP2.getBlock() && IP1.getPoint() == IP2.getPoint();
 }
 
+/// This is wrapper over IRBuilderBase::restoreIP that also restores the current
+/// debug location to the last instruction in the specified basic block if the
+/// insert point points to the end of the block.
+static void restoreIPandDebugLoc(llvm::IRBuilderBase &Builder,
+                                 llvm::IRBuilderBase::InsertPoint IP) {
+  Builder.restoreIP(IP);
+  llvm::BasicBlock *BB = Builder.GetInsertBlock();
+  llvm::BasicBlock::iterator I = Builder.GetInsertPoint();
+  if (!BB->empty() && I == BB->end())
+    Builder.SetCurrentDebugLocation(BB->back().getStableDebugLoc());
+}
+
 static bool isValidWorkshareLoopScheduleType(OMPScheduleType SchedType) {
   // Valid ordered/unordered and base algorithm combinations.
   switch (SchedType & ~OMPScheduleType::MonotonicityMask) {
@@ -8993,7 +9005,7 @@ Error OpenMPIRBuilder::emitOffloadingArrays(
     ArrayType *SizeArrayType = ArrayType::get(Int64Ty, Info.NumberOfPtrs);
     Info.RTArgs.SizesArray = Builder.CreateAlloca(
         SizeArrayType, /* ArraySize = */ nullptr, ".offload_sizes");
-    Builder.restoreIP(CodeGenIP);
+    restoreIPandDebugLoc(Builder, CodeGenIP);
   } else {
     auto *SizesArrayInit = ConstantArray::get(
         ArrayType::get(Int64Ty, ConstSizes.size()), ConstSizes);
@@ -9012,7 +9024,7 @@ Error OpenMPIRBuilder::emitOffloadingArrays(
       AllocaInst *Buffer = Builder.CreateAlloca(
           SizeArrayType, /* ArraySize = */ nullptr, ".offload_sizes");
       Buffer->setAlignment(OffloadSizeAlign);
-      Builder.restoreIP(CodeGenIP);
+      restoreIPandDebugLoc(Builder, CodeGenIP);
       Builder.CreateMemCpy(
           Buffer, M.getDataLayout().getPrefTypeAlign(Buffer->getType()),
           SizesArrayGbl, OffloadSizeAlign,
@@ -9022,7 +9034,7 @@ Error OpenMPIRBuilder::emitOffloadingArrays(
 
       Info.RTArgs.SizesArray = Buffer;
     }
-    Builder.restoreIP(CodeGenIP);
+    restoreIPandDebugLoc(Builder, CodeGenIP);
   }
 
   // The map types are always constant so we don't need to generate code to

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

@@ -4356,9 +4356,11 @@ createAlteredByCaptureMap(MapInfoData &mapData,
 
         if (!isPtrTy) {
           auto curInsert = builder.saveIP();
+          llvm::DebugLoc DbgLoc = builder.getCurrentDebugLocation();
           builder.restoreIP(findAllocaInsertPoint(builder, moduleTranslation));
           auto *memTempAlloc =
               builder.CreateAlloca(builder.getPtrTy(), nullptr, ".casted");
+          builder.SetCurrentDebugLocation(DbgLoc);
           builder.restoreIP(curInsert);
 
           builder.CreateStore(newV, memTempAlloc);

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

@@ -0,0 +1,45 @@
+// RUN: mlir-translate -mlir-to-llvmir %s
+
+module attributes {llvm.target_triple = "x86_64-unknown-linux-gnu", omp.is_gpu = false, omp.is_target_device = false, omp.target_triples = ["amdgcn-amd-amdhsa"]} {
+  omp.private {type = private} @_QFFfnEv_private_i32 : i32 loc(#loc1)
+  llvm.func internal @_QFPfn() {
+    %0 = llvm.mlir.constant(1 : i64) : i64 loc(#loc1)
+    %1 = llvm.alloca %0 x i32 {bindc_name = "v"} : (i64) -> !llvm.ptr loc(#loc1)
+    %2 = llvm.mlir.constant(1 : i32) : i32
+    omp.parallel private(@_QFFfnEv_private_i32 %1 -> %arg0 : !llvm.ptr) {
+      llvm.store %2, %arg0 : i32, !llvm.ptr loc(#loc2)
+      %4 = omp.map.info var_ptr(%arg0 : !llvm.ptr, i32) map_clauses(implicit, exit_release_or_enter_alloc) capture(ByCopy) -> !llvm.ptr {name = "v"} loc(#loc2)
+      omp.target map_entries(%4 -> %arg1 : !llvm.ptr) {
+        %5 = llvm.mlir.constant(1 : i32) : i32
+        %6 = llvm.load %arg1 : !llvm.ptr -> i32 loc(#loc3)
+        %7 = llvm.add %6, %5 : i32 loc(#loc3)
+        llvm.store %7, %arg1 : i32, !llvm.ptr loc(#loc3)
+        omp.terminator loc(#loc3)
+      } loc(#loc7)
+      omp.terminator
+    } loc(#loc4)
+    llvm.return
+  } loc(#loc6)
+}
+
+#di_file = #llvm.di_file<"target.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_program, types = #di_null_type>
+#di_subprogram = #llvm.di_subprogram<id = distinct[1]<>,
+  compileUnit = #di_compile_unit, scope = #di_file, name = "main",
+  file = #di_file, subprogramFlags = "Definition|MainSubprogram",
+  type = #di_subroutine_type>
+#di_subprogram1 = #llvm.di_subprogram<compileUnit = #di_compile_unit,
+  name = "target", file = #di_file, subprogramFlags = "Definition",
+  type = #di_subroutine_type>
+
+#loc1 = loc("test.f90":7:15)
+#loc2 = loc("test.f90":1:7)
+#loc3 = loc("test.f90":3:7)
+#loc4 = loc("test.f90":16:7)
+#loc6 = loc(fused<#di_subprogram>[#loc1])
+#loc7 = loc(fused<#di_subprogram1>[#loc3])