diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h index 5bbaa8c208b8c..54908d01bc006 100644 --- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h +++ b/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_" + /// + ".var" for "omp critical" directives; 2) diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp index d65ed8c11d86c..e634e221a96de 100644 --- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp +++ b/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 &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"); } diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp index bef227f2c583e..ab836251e15eb 100644 --- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp +++ b/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(); } diff --git a/mlir/test/Target/LLVMIR/openmp-task-target-device.mlir b/mlir/test/Target/LLVMIR/openmp-task-target-device.mlir new file mode 100644 index 0000000000000..c167604bcd12c --- /dev/null +++ b/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} { + %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 + } +}