Fix incompatible DI location op types

I didn't notice in the initial review, but the type in the expression
should match the type of the location op (or at least be a compatible
size), and in a merge from amd-debug the verifier now catches this
issue.

I also didn't realize I had missed the flang/openmp changes when
creating amd-debug. My plan is to land the fix here in amd-staging,
then complete the merge from amd-debug to amd-staging, and then
port the code from amd-staging to amd-debug.

Author: slinder1

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

@@ -5739,21 +5739,23 @@ static void addAllocasForDeclareTargetFunctionPointerArgs(
   // below. The users are updated accordingly.
   for (auto &Arg : Fn->args()) {
     if (Arg.getType()->isPointerTy()) {
-      llvm::Value *V = builder.CreateAlloca(Arg.getType(), allocaAS, nullptr);
-      if (allocaAS != defaultAS)
-        V = ompBuilder->Builder.CreateAddrSpaceCast(
-            V, builder.getPtrTy(defaultAS));
-      llvm::StoreInst *Store = builder.CreateStore(&Arg, V);
-      llvm::Value *Load = builder.CreateLoad(Arg.getType(), V);
+      llvm::Value *AllocaV =
+          builder.CreateAlloca(Arg.getType(), allocaAS, nullptr);
+      llvm::Value *GenericV = allocaAS == defaultAS
+                                  ? AllocaV
+                                  : ompBuilder->Builder.CreateAddrSpaceCast(
+                                        AllocaV, builder.getPtrTy(defaultAS));
+      llvm::StoreInst *Store = builder.CreateStore(&Arg, GenericV);
+      llvm::Value *Load = builder.CreateLoad(Arg.getType(), GenericV);
       llvm::SmallVector<llvm::DbgVariableIntrinsic *> DbgUsers;
       llvm::SmallVector<llvm::DbgVariableRecord *> DPUsers;
       llvm::findDbgUsers(DbgUsers, &Arg, &DPUsers);
       for (auto *DVI : DbgUsers) {
-        DVI->replaceVariableLocationOp(&Arg, V);
+        DVI->replaceVariableLocationOp(&Arg, AllocaV);
         DVI->setExpression(Expr);
       }
       for (auto *DVR : DPUsers) {
-        DVR->replaceVariableLocationOp(&Arg, V);
+        DVR->replaceVariableLocationOp(&Arg, AllocaV);
         DVR->setExpression(Expr);
       }
       Arg.replaceUsesWithIf(Load, [&](const llvm::Use &U) -> bool {
@@ -5783,16 +5785,22 @@ static void updateDebugInfoForDeclareTargetFunctions(
     // Skip if an expression is already present.
     if ((Old != nullptr) && (Old->getNumElements() != 0))
       return;
-    for (auto Loc : DR->location_ops()) {
-      llvm::Type *Ty = Loc->getType();
-      if (auto *Ref = dyn_cast<llvm::AddrSpaceCastInst>(Loc))
-        Ty = Ref->getPointerOperand()->getType();
-      llvm::DIExprBuilder EB(Fn->getContext());
-      EB.append<llvm::DIOp::Arg>(0u, Ty);
-      EB.append<llvm::DIOp::Deref>(Loc->getType());
-      DR->setExpression(EB.intoExpression());
-      break;
+    // Skip if the there are multiple inputs.
+    // FIXME: Could this be an assert? More to the point, can we do this at the
+    // point of generating the intrinsics to begin with, rather than fixing them
+    // up here?
+    if (DR->getNumVariableLocationOps() != 1u)
+      return;
+    auto Loc = DR->getVariableLocationOp(0u);
+    llvm::Type *Ty = Loc->getType();
+    if (auto *Ref = dyn_cast<llvm::AddrSpaceCastInst>(Loc)) {
+      DR->replaceVariableLocationOp(0u, Loc->stripPointerCasts());
+      Ty = Ref->getPointerOperand()->getType();
     }
+    llvm::DIExprBuilder EB(Fn->getContext());
+    EB.append<llvm::DIOp::Arg>(0u, Ty);
+    EB.append<llvm::DIOp::Deref>(Loc->getType());
+    DR->setExpression(EB.intoExpression());
   };
 
   for (llvm::Instruction &I : instructions(Fn)) {

mlir/test/Target/LLVMIR/omptarget-decl-target-fn-debug-amdgpu.mlir

@@ -27,7 +27,7 @@ module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_target_devic
 // CHECK:   %[[CAST:[0-9]+]] = addrspacecast ptr addrspace(5) %[[AL]]
 // CHECK:   store ptr %[[ARG]], ptr %[[CAST]]{{.*}}
 // CHECK:   load ptr, ptr %[[CAST]]{{.*}}
-// CHECK: #dbg_declare(ptr %[[CAST]], ![[A:[0-9]+]], !DIExpression(DIOpArg(0, ptr addrspace(5)), DIOpDeref(ptr), DIOpDeref(ptr)), !{{.*}})
+// CHECK: #dbg_declare(ptr addrspace(5) %[[AL]], ![[A:[0-9]+]], !DIExpression(DIOpArg(0, ptr addrspace(5)), DIOpDeref(ptr), DIOpDeref(ptr)), !{{.*}})
 // CHECK: }
 // CHECK: ![[SP]] = {{.*}}!DISubprogram(name: "add"{{.*}})
 // CHECK: ![[A]] = !DILocalVariable(name: "a", arg: 1, scope: ![[SP]]{{.*}})