[mlir][openmp] - Fix crash in OpenMPIRBuilder when converting to LLVMIR

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

@@ -465,6 +465,10 @@ class OpenMPIRBuilder {
 
   void setConfig(OpenMPIRBuilderConfig C) { Config = C; }
 
+  /// Remove all references or state about Func that OpenMPIRBuilder may
+  /// be keeping
+  void dropFunction(Function *Func);
+
   /// Finalize the underlying module, e.g., by outlining regions.
   /// \param Fn                    The function to be finalized. If not used,
   ///                              all functions are finalized.
@@ -1518,6 +1522,9 @@ class OpenMPIRBuilder {
   /// Add a new region that will be outlined later.
   void addOutlineInfo(OutlineInfo &&OI) { OutlineInfos.emplace_back(OI); }
 
+  /// Remove outlining information if it refers to a certain function
+  void removeFuncFromOutlineInfo(llvm::Function *Func);
+
   /// An ordered map of auto-generated variables to their unique names.
   /// It stores variables with the following names: 1) ".gomp_critical_user_" +
   /// <critical_section_name> + ".var" for "omp critical" directives; 2)

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -488,6 +488,14 @@ void OpenMPIRBuilderConfig::setHasRequiresDynamicAllocators(bool Value) {
 // OpenMPIRBuilder
 //===----------------------------------------------------------------------===//
 
+void OpenMPIRBuilder::removeFuncFromOutlineInfo(llvm::Function *Func) {
+  OutlineInfos.erase(std::remove_if(OutlineInfos.begin(), OutlineInfos.end(),
+                                    [&](const OutlineInfo &OI) {
+                                      return OI.getFunction() == Func;
+                                    }),
+                     OutlineInfos.end());
+}
+
 void OpenMPIRBuilder::getKernelArgsVector(TargetKernelArgs &KernelArgs,
                                           IRBuilderBase &Builder,
                                           SmallVector<Value *> &ArgsVector) {
@@ -794,6 +802,10 @@ void OpenMPIRBuilder::finalize(Function *Fn) {
     createOffloadEntriesAndInfoMetadata(ErrorReportFn);
 }
 
+void OpenMPIRBuilder::dropFunction(Function *Func) {
+  removeFuncFromOutlineInfo(Func);
+}
+
 OpenMPIRBuilder::~OpenMPIRBuilder() {
   assert(OutlineInfos.empty() && "There must be no outstanding outlinings");
 }

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

@@ -2853,6 +2853,8 @@ convertDeclareTargetAttr(Operation *op, mlir::omp::DeclareTargetAttr attribute,
       if (declareType == omp::DeclareTargetDeviceType::host) {
         llvm::Function *llvmFunc =
             moduleTranslation.lookupFunction(funcOp.getName());
+
+        moduleTranslation.getOpenMPBuilder()->dropFunction(llvmFunc);
         llvmFunc->dropAllReferences();
         llvmFunc->eraseFromParent();
       }

mlir/test/Target/LLVMIR/openmp-task-target-device.mlir

@@ -0,0 +1,27 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// This tests the fix for https://github.com/llvm/llvm-project/issues/84606
+// We are only interested in ensuring that the -mlir-to-llmvir pass doesn't crash.
+// CHECK: {{.*}} = add i32 {{.*}}, 5
+module attributes {omp.is_target_device = true } {
+  llvm.func @_QQmain() attributes {fir.bindc_name = "main", omp.declare_target = #omp.declaretarget<device_type = (host), capture_clause = (to)>} {
+    %0 = llvm.mlir.constant(0 : i32) : i32
+    %1 = llvm.mlir.constant(1 : i64) : i64
+    %2 = llvm.alloca %1 x i32 {bindc_name = "a"} : (i64) -> !llvm.ptr<5>
+    %3 = llvm.addrspacecast %2 : !llvm.ptr<5> to !llvm.ptr
+    omp.task {
+      llvm.store %0, %3 : i32, !llvm.ptr
+      omp.terminator
+    }
+    %4 = omp.map_info var_ptr(%3 : !llvm.ptr, i32) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = "a"}
+    omp.target map_entries(%4 -> %arg0 : !llvm.ptr) {
+    ^bb0(%arg0: !llvm.ptr):
+      %5 = llvm.mlir.constant(5 : i32) : i32
+      %6 = llvm.load %arg0  : !llvm.ptr -> i32
+      %7 = llvm.add %6, %5  : i32
+      llvm.store %7, %arg0  : i32, !llvm.ptr
+      omp.terminator
+    }
+    llvm.return
+  }
+}