From b03cc7b63d62485405d43998986aeb94de785b98 Mon Sep 17 00:00:00 2001
From: Pranav Bhandarkar <pranav.bhandarkar@amd.com>
Date: Fri, 8 Mar 2024 10:56:05 -0600
Subject: [PATCH 1/2] [mlir][openmp] - Fix crash in OpenMPIRBuilder when
converting to LLVMIR
This patch fixes an issue in the conversion of the omp MLIR dialect
to LLVMIR using the OpenMPIRBuilder seen during offloading.
The finalize method in OpenMPIRBuilder outlines regions marked for
outlining. To do so, it uses a datastructure that holds blocks that
need to be outline. However, if these regions belong to host code in
a device module then these regions will have been removed when the
"omp.declare_target" attribute is converted. As a result the finalize
method ends up accessing bad data and crashes.
---
.../llvm/Frontend/OpenMP/OMPIRBuilder.h | 6 +++++
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 12 +++++++++
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 2 ++
.../LLVMIR/openmp-task-target-device.mlir | 27 +++++++++++++++++++
4 files changed, 47 insertions(+)
create mode 100644 mlir/test/Target/LLVMIR/openmp-task-target-device.mlir
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index 5bbaa8c208b8c..2b83a1b91e3b2 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -465,6 +465,9 @@ 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 +1521,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)
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<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");
}
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<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
+ }
+}
From 811c5004264ab9ac3df2926af8848b171e3baad1 Mon Sep 17 00:00:00 2001
From: Pranav Bhandarkar <pranav.bhandarkar@amd.com>
Date: Sat, 9 Mar 2024 01:30:39 -0600
Subject: [PATCH 2/2] Formating changes in OMPIRBuilder.h
---
llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h | 7 ++++---
1 file changed, 4 insertions(+), 3 deletions(-)
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index 2b83a1b91e3b2..54908d01bc006 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -465,9 +465,10 @@ class OpenMPIRBuilder {
void setConfig(OpenMPIRBuilderConfig C) { Config = C; }
- // Remove all references or state about Func that OpenMPIRBuilder may
- // be keeping
+ /// 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.
@@ -1521,7 +1522,7 @@ 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
+ /// 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.