Add allocas for arguments passed in by ref.

abidh · abidh · commit fec5377d8a98 · 2025-03-31T12:45:34.000+01:00
Initially the PR was adding alloca because of the problem with using Argument as location. Although that problem has been resolved and I tried to use them, the debug info that comes out is very poor and mostly ususable.

This commit now adds alloca and all the case so that code is same with and without debug information enabled.
diff --git a/flang/test/Integration/amdgpu/debug-declare-target-function-var.f90 b/flang/test/Integration/amdgpu/debug-declare-target-function-var.f90
@@ -13,8 +13,8 @@ function add(a, b) result(ret)
 end
 
 !CHECK: define float @add_({{.*}}){{.*}}!dbg ![[SP:[0-9]+]] {
-!CHECK: #dbg_declare({{.*}}, ![[A:[0-9]+]], !DIExpression(DIOpArg(0, ptr), DIOpDeref(ptr)), !{{.*}})
-!CHECK: #dbg_declare({{.*}}, ![[B:[0-9]+]], !DIExpression(DIOpArg(0, ptr), DIOpDeref(ptr)), !{{.*}})
+!CHECK: #dbg_declare({{.*}}, ![[A:[0-9]+]], !DIExpression(DIOpArg(0, ptr addrspace(5)), DIOpDeref(ptr), DIOpDeref(ptr)), !{{.*}})
+!CHECK: #dbg_declare({{.*}}, ![[B:[0-9]+]], !DIExpression(DIOpArg(0, ptr addrspace(5)), DIOpDeref(ptr), DIOpDeref(ptr)), !{{.*}})
 !CHECK: #dbg_declare({{.*}}, ![[RET:[0-9]+]], !DIExpression(DIOpArg(0, ptr addrspace(5)), DIOpDeref(ptr)), !{{.*}})
 !CHECK: }
 !CHECK: ![[SP]] = {{.*}}!DISubprogram(name: "add"{{.*}})
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -31,16 +31,17 @@
 #include "llvm/Frontend/OpenMP/OMPDeviceConstants.h"
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/ReplaceConstant.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
-#include "llvm/IR/InstIterator.h"
-#include "llvm/IR/IntrinsicInst.h"
 
 #include <any>
 #include <cstdint>
@@ -5464,8 +5465,10 @@ static void updateDebugInfoForDeclareTargetVariables(
   }
 }
 
-// This function adds DIOp based expressions in declare target function to
-// generate valid debug info for AMDGPU.
+// This function handle any adjustments needed in declare target function to
+// generate valid debug info for AMDGPU. It does 2 main things:
+// 1. Add alloca for arguments passed by reference.
+// 2. Add DIOp based expressions
 
 static void updateDebugInfoForDeclareTargetFunctions(
     llvm::Function *Fn, LLVM::ModuleTranslation &moduleTranslation) {
@@ -5475,23 +5478,80 @@ static void updateDebugInfoForDeclareTargetFunctions(
   if (!llvm::Triple(M.getTargetTriple()).isAMDGPU())
     return;
 
-  auto UpdateDebugRecord = [&](auto *DR) {
+  llvm::IRBuilderBase &builder = ompBuilder->Builder;
+  llvm::OpenMPIRBuilder::InsertPointTy curInsert = builder.saveIP();
+  unsigned int allocaAS = M.getDataLayout().getAllocaAddrSpace();
+  unsigned int defaultAS = M.getDataLayout().getProgramAddressSpace();
+
+  builder.SetInsertPoint(Fn->getEntryBlock().getFirstInsertionPt());
+
+  llvm::Type *PtrTy = builder.getPtrTy(defaultAS);
+  llvm::Type *AllocaPtrTy = builder.getPtrTy(allocaAS);
+  llvm::DIExprBuilder EB(Fn->getContext());
+  EB.append<llvm::DIOp::Arg>(0u, AllocaPtrTy);
+  EB.append<llvm::DIOp::Deref>(PtrTy);
+  EB.append<llvm::DIOp::Deref>(PtrTy);
+  llvm::DIExpression *Expr = EB.intoExpression();
+
+  // flang does not generate allocas for the arguments that are passed by ref.
+  // When the Argument is the location, the quality of the debug information is
+  // poor. The variables are defines on very few addresses and show up as
+  // optimized in most places. One of the reason is the interaction of DI-Op
+  // based ops and regular ones.
+  // Generating alloca seems like the best thing which is done in the loop
+  // 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::SmallVector<llvm::DbgVariableIntrinsic *> DbgUsers;
+      llvm::SmallVector<llvm::DbgVariableRecord *> DPUsers;
+      llvm::findDbgUsers(DbgUsers, &Arg, &DPUsers);
+      for (auto *DVI : DbgUsers) {
+        DVI->replaceVariableLocationOp(&Arg, V);
+        DVI->setExpression(Expr);
+      }
+      for (auto *DVR : DPUsers) {
+        DVR->replaceVariableLocationOp(&Arg, V);
+        DVR->setExpression(Expr);
+      }
+      Arg.replaceUsesWithIf(Load, [&](const llvm::Use &U) -> bool {
+        // We dont want to replace Arg from the store we created above.
+        if (const auto *SI = dyn_cast<llvm::StoreInst>(U.getUser()))
+          return SI != Store;
+        return true;
+      });
+    }
+  }
+  builder.restoreIP(curInsert);
+
+  auto AddExpression = [&](auto *DR) {
+    llvm::DIExpression *Old = DR->getExpression();
+    // Skip if an expression is already present.
+    if ((Old != nullptr) && (Old->getNumElements() != 0))
+      return;
     for (auto Loc : DR->location_ops()) {
-      llvm::DIExprBuilder ExprBuilder(Fn->getContext());
       llvm::Type *Ty = Loc->getType();
       if (auto *Ref = dyn_cast<llvm::AddrSpaceCastInst>(Loc))
         Ty = Ref->getPointerOperand()->getType();
-      ExprBuilder.append<llvm::DIOp::Arg>(0u, Ty);
-      ExprBuilder.append<llvm::DIOp::Deref>(Loc->getType());
-      DR->setExpression(ExprBuilder.intoExpression());
+      llvm::DIExprBuilder EB(Fn->getContext());
+      EB.append<llvm::DIOp::Arg>(0u, Ty);
+      EB.append<llvm::DIOp::Deref>(Loc->getType());
+      DR->setExpression(EB.intoExpression());
+      break;
     }
   };
+
   for (llvm::Instruction &I : instructions(Fn)) {
     if (auto *DDI = dyn_cast<llvm::DbgVariableIntrinsic>(&I))
-      UpdateDebugRecord(DDI);
+      AddExpression(DDI);
 
     for (llvm::DbgVariableRecord &DVR : filterDbgVars(I.getDbgRecordRange()))
-      UpdateDebugRecord(&DVR);
+      AddExpression(&DVR);
   }
 }
 
diff --git a/mlir/test/Target/LLVMIR/omptarget-decl-target-fn-debug-amdgpu.mlir b/mlir/test/Target/LLVMIR/omptarget-decl-target-fn-debug-amdgpu.mlir
@@ -23,7 +23,11 @@ module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_target_devic
 #loc3 = loc(fused<#sp>[#loc1])
 
 // CHECK: define{{.*}}@add(ptr %[[ARG:[0-9]+]]){{.*}}!dbg ![[SP:[0-9]+]] {
-// CHECK: #dbg_declare(ptr %[[ARG]], ![[A:[0-9]+]], !DIExpression(DIOpArg(0, ptr), DIOpDeref(ptr)), !{{.*}})
+// CHECK:   %[[AL:[0-9]+]] = alloca{{.*}}
+// 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: }
 // CHECK: ![[SP]] = {{.*}}!DISubprogram(name: "add"{{.*}})
 // CHECK: ![[A]] = !DILocalVariable(name: "a", arg: 1, scope: ![[SP]]{{.*}})
diff --git a/mlir/test/Target/LLVMIR/omptarget-wsloop.mlir b/mlir/test/Target/LLVMIR/omptarget-wsloop.mlir
@@ -38,8 +38,12 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memo
 // CHECK: define void @[[FUNC0:.*]](ptr %[[ARG0:.*]])
 // CHECK:   %[[STRUCTARG:.*]] = alloca { ptr }, align 8, addrspace(5)
 // CHECK:   %[[STRUCTARG_ASCAST:.*]] = addrspacecast ptr addrspace(5) %[[STRUCTARG]] to ptr
+// CHECK:   %[[AL:[0-9]+]] = alloca{{.*}}
+// CHECK:   %[[CAST:[0-9]+]] = addrspacecast ptr addrspace(5) %[[AL]]
+// CHECK:   store ptr %[[ARG0]], ptr %[[CAST]]{{.*}}
+// CHECK:   %[[LOAD:[0-9]+]] = load ptr, ptr %[[CAST]]{{.*}}
 // CHECK:   %[[GEP:.*]] = getelementptr { ptr }, ptr %[[STRUCTARG_ASCAST]], i32 0, i32 0
-// CHECK:   store ptr %[[ARG0]], ptr %[[GEP]], align 8
+// CHECK:   store ptr %[[LOAD]], ptr %[[GEP]], align 8
 // CHECK:   %[[NUM_THREADS:.*]] = call i32 @omp_get_num_threads()
 // CHECK:   call void @__kmpc_for_static_loop_4u(ptr addrspacecast (ptr addrspace(1) @[[GLOB1:[0-9]+]] to ptr), ptr @[[LOOP_BODY_FN:.*]], ptr %[[STRUCTARG_ASCAST]], i32 9, i32 %[[NUM_THREADS]], i32 0)
 
