Remove extra allocas from declare target functions. (llvm#3777)

Author: ronlieb

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

@@ -5896,71 +5896,6 @@ static void updateDebugInfoForDeclareTargetVariables(
   }
 }
 
-static void addAllocasForDeclareTargetFunctionPointerArgs(
-    llvm::Function *Fn, LLVM::ModuleTranslation &moduleTranslation) {
-  llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
-  llvm::Module &M = ompBuilder->M;
-
-  if (!llvm::Triple(M.getTargetTriple()).isAMDGPU())
-    return;
-
-  if (Fn->empty())
-    return;
-
-  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 *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(&Arg, DPUsers);
-      for (auto *DVI : DbgUsers) {
-        DVI->replaceVariableLocationOp(&Arg, AllocaV);
-        DVI->setExpression(Expr);
-      }
-      for (auto *DVR : DPUsers) {
-        DVR->replaceVariableLocationOp(&Arg, AllocaV);
-        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);
-}
-
 // This function Add DIOp based expressions to the debug records in the
 // declare target functions.
 
@@ -5984,13 +5919,16 @@ static void updateDebugInfoForDeclareTargetFunctions(
     if (DR->getNumVariableLocationOps() != 1u)
       return;
     auto Loc = DR->getVariableLocationOp(0u);
-    if (!isa<llvm::AllocaInst>(Loc->stripPointerCasts()))
-      return;
-    llvm::AllocaInst *AI = cast<llvm::AllocaInst>(Loc->stripPointerCasts());
-    DR->replaceVariableLocationOp(0u, AI);
     llvm::DIExprBuilder EB(Fn->getContext());
-    EB.append<llvm::DIOp::Arg>(0u, AI->getType());
-    EB.append<llvm::DIOp::Deref>(AI->getAllocatedType());
+    if (auto AI = dyn_cast<llvm::AllocaInst>(Loc->stripPointerCasts())) {
+      DR->replaceVariableLocationOp(0u, AI);
+      EB.append<llvm::DIOp::Arg>(0u, AI->getType());
+      EB.append<llvm::DIOp::Deref>(AI->getAllocatedType());
+    } else if (Loc->getType()->isPointerTy()) {
+      EB.append<llvm::DIOp::Arg>(0u, Loc->getType());
+      EB.append<llvm::DIOp::Deref>(Loc->getType());
+    } else
+      EB.append<llvm::DIOp::Arg>(0u, Loc->getType());
     DR->setExpression(EB.intoExpression());
   };
 
@@ -6027,11 +5965,8 @@ convertDeclareTargetAttr(Operation *op, mlir::omp::DeclareTargetAttr attribute,
       if (declareType == omp::DeclareTargetDeviceType::host) {
         llvmFunc->dropAllReferences();
         llvmFunc->eraseFromParent();
-      } else {
-        addAllocasForDeclareTargetFunctionPointerArgs(llvmFunc,
-                                                      moduleTranslation);
+      } else
         updateDebugInfoForDeclareTargetFunctions(llvmFunc, moduleTranslation);
-      }
     }
     return success();
   }

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

@@ -1,5 +1,5 @@
 // RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
-// XFAIL: *
+
 #file = #llvm.di_file<"target.f90" in "">
 #cu = #llvm.di_compile_unit<id = distinct[0]<>,
  sourceLanguage = DW_LANG_Fortran95, file = #file, isOptimized = false,
@@ -23,11 +23,7 @@ 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:   %[[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 addrspace(5) %[[AL]], ![[A:[0-9]+]], !DIExpression(DIOpArg(0, ptr addrspace(5)), DIOpDeref(ptr), DIOpDeref(ptr)), !{{.*}})
+// CHECK: #dbg_declare(ptr %[[ARG]], ![[A:[0-9]+]], !DIExpression(DIOpArg(0, ptr), DIOpDeref(ptr)), !{{.*}})
 // CHECK: }
 // CHECK: ![[SP]] = {{.*}}!DISubprogram(name: "add"{{.*}})
 // CHECK: ![[A]] = !DILocalVariable(name: "a", arg: 1, scope: ![[SP]]{{.*}})