[flang-commits] [flang] [llvm] [mlir] [mlir][OpenMP] Add the tgt_entry_info attribute (PR #78223)
Fabian Mora via flang-commits
flang-commits at lists.llvm.org
Mon Jan 15 18:24:49 PST 2024
https://github.com/fabianmcg created https://github.com/llvm/llvm-project/pull/78223
This patch adds the `omp.tgt_entry_info` attribute. This attribute provides
information to identify offload entries uniquely and partially reflects the
information in the `llvm::TargetRegionEntryInfo` struct.
An `info` parameter was added to `TargetOp` to specify the offload entry
information from the operation explicitly. Both the host and device versions
must have the same `info` attribute; otherwise, the constructs won't correctly
map between each other.
This patch is required to enable JIT compilation for the OMP dialect, as
the entry array has to be fully constructed in the IR instead of using
sections to implicitly construct it.
Note: Ignore the base commits.
>From 96ca7efc81c85ee6011add476a5987f6e0efacef Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Mon, 15 Jan 2024 15:01:42 +0000
Subject: [PATCH 1/3] [llvm][OpenMPIRBuilder] Allow to not register offload
entries in the entry manager This patch adds an optional field in the create
target method to store the offload entry in a custom location and not
register the entry in the entry manager. This change is required to enable
JIT compilation in MLIR for OpenMP target offload ops, as arrays of entries
are handled differently for standalone MLIR compilation.
---
.../llvm/Frontend/OpenMP/OMPIRBuilder.h | 53 ++++++++++++-------
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 36 ++++++++-----
2 files changed, 57 insertions(+), 32 deletions(-)
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index 669104307fa0e2..a584552d15954d 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -298,9 +298,10 @@ class OffloadEntriesInfoManager {
void initializeTargetRegionEntryInfo(const TargetRegionEntryInfo &EntryInfo,
unsigned Order);
/// Register target region entry.
- void registerTargetRegionEntryInfo(TargetRegionEntryInfo EntryInfo,
- Constant *Addr, Constant *ID,
- OMPTargetRegionEntryKind Flags);
+ void registerTargetRegionEntryInfo(
+ TargetRegionEntryInfo EntryInfo, Constant *Addr, Constant *ID,
+ OMPTargetRegionEntryKind Flags,
+ OffloadEntryInfoTargetRegion *EntryInfoStorage = nullptr);
/// Return true if a target region entry with the provided information
/// exists.
bool hasTargetRegionEntryInfo(TargetRegionEntryInfo EntryInfo,
@@ -2139,10 +2140,14 @@ class OpenMPIRBuilder {
/// \param GenerateFunctionCallback The callback function to generate the code
/// \param OutlinedFunction Pointer to the outlined function
/// \param EntryFnIDName Name of the ID o be created
- void emitTargetRegionFunction(TargetRegionEntryInfo &EntryInfo,
- FunctionGenCallback &GenerateFunctionCallback,
- bool IsOffloadEntry, Function *&OutlinedFn,
- Constant *&OutlinedFnID);
+ /// \param EntryInfoStorage Optional storage location for the offload target
+ /// region entry.
+ void emitTargetRegionFunction(
+ TargetRegionEntryInfo &EntryInfo,
+ FunctionGenCallback &GenerateFunctionCallback, bool IsOffloadEntry,
+ Function *&OutlinedFn, Constant *&OutlinedFnID,
+ OffloadEntriesInfoManager::OffloadEntryInfoTargetRegion
+ *EntryInfoStorage = nullptr);
/// Registers the given function and sets up the attribtues of the function
/// Returns the FunctionID.
@@ -2152,10 +2157,13 @@ class OpenMPIRBuilder {
/// \param OutlinedFunction Pointer to the outlined function
/// \param EntryFnName Name of the outlined function
/// \param EntryFnIDName Name of the ID o be created
- Constant *registerTargetRegionFunction(TargetRegionEntryInfo &EntryInfo,
- Function *OutlinedFunction,
- StringRef EntryFnName,
- StringRef EntryFnIDName);
+ /// \param EntryInfoStorage Optional storage location for the offload target
+ /// region entry.
+ Constant *registerTargetRegionFunction(
+ TargetRegionEntryInfo &EntryInfo, Function *OutlinedFunction,
+ StringRef EntryFnName, StringRef EntryFnIDName,
+ OffloadEntriesInfoManager::OffloadEntryInfoTargetRegion
+ *EntryInfoStorage = nullptr);
/// Type of BodyGen to use for region codegen
///
@@ -2225,15 +2233,20 @@ class OpenMPIRBuilder {
/// \param BodyGenCB Callback that will generate the region code.
/// \param ArgAccessorFuncCB Callback that will generate accessors
/// instructions for passed in target arguments where neccessary
- InsertPointTy createTarget(const LocationDescription &Loc,
- OpenMPIRBuilder::InsertPointTy AllocaIP,
- OpenMPIRBuilder::InsertPointTy CodeGenIP,
- TargetRegionEntryInfo &EntryInfo, int32_t NumTeams,
- int32_t NumThreads,
- SmallVectorImpl<Value *> &Inputs,
- GenMapInfoCallbackTy GenMapInfoCB,
- TargetBodyGenCallbackTy BodyGenCB,
- TargetGenArgAccessorsCallbackTy ArgAccessorFuncCB);
+ /// \param EntryInfoStorage Optional storage location for the offload target
+ /// region entry. If the pointer is null the entry gets registered in the
+ /// `OffloadEntriesInfoManager`, otherwise it gets stored in
+ /// `EntryInfoStorage` and doesn't get registered in
+ /// `OffloadEntriesInfoManager`.
+ InsertPointTy createTarget(
+ const LocationDescription &Loc, OpenMPIRBuilder::InsertPointTy AllocaIP,
+ OpenMPIRBuilder::InsertPointTy CodeGenIP,
+ TargetRegionEntryInfo &EntryInfo, int32_t NumTeams, int32_t NumThreads,
+ SmallVectorImpl<Value *> &Inputs, GenMapInfoCallbackTy GenMapInfoCB,
+ TargetBodyGenCallbackTy BodyGenCB,
+ TargetGenArgAccessorsCallbackTy ArgAccessorFuncCB,
+ OffloadEntriesInfoManager::OffloadEntryInfoTargetRegion
+ *EntryInfoStorage = nullptr);
/// Returns __kmpc_for_static_init_* runtime function for the specified
/// size \a IVSize and sign \a IVSigned. Will create a distribute call
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index f6cf358119fb71..c5498d6c67b455 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -4741,7 +4741,8 @@ Constant *OpenMPIRBuilder::createTargetRegionEntryAddr(Function *OutlinedFn,
void OpenMPIRBuilder::emitTargetRegionFunction(
TargetRegionEntryInfo &EntryInfo,
FunctionGenCallback &GenerateFunctionCallback, bool IsOffloadEntry,
- Function *&OutlinedFn, Constant *&OutlinedFnID) {
+ Function *&OutlinedFn, Constant *&OutlinedFnID,
+ OffloadEntriesInfoManager::OffloadEntryInfoTargetRegion *EntryInfoStorage) {
SmallString<64> EntryFnName;
OffloadInfoManager.getTargetRegionEntryFnName(EntryFnName, EntryInfo);
@@ -4761,20 +4762,22 @@ void OpenMPIRBuilder::emitTargetRegionFunction(
? std::string(EntryFnName)
: createPlatformSpecificName({EntryFnName, "region_id"});
- OutlinedFnID = registerTargetRegionFunction(EntryInfo, OutlinedFn,
- EntryFnName, EntryFnIDName);
+ OutlinedFnID = registerTargetRegionFunction(
+ EntryInfo, OutlinedFn, EntryFnName, EntryFnIDName, EntryInfoStorage);
}
Constant *OpenMPIRBuilder::registerTargetRegionFunction(
TargetRegionEntryInfo &EntryInfo, Function *OutlinedFn,
- StringRef EntryFnName, StringRef EntryFnIDName) {
+ StringRef EntryFnName, StringRef EntryFnIDName,
+ OffloadEntriesInfoManager::OffloadEntryInfoTargetRegion *EntryInfoStorage) {
if (OutlinedFn)
setOutlinedTargetRegionFunctionAttributes(OutlinedFn);
auto OutlinedFnID = createOutlinedFunctionID(OutlinedFn, EntryFnIDName);
auto EntryAddr = createTargetRegionEntryAddr(OutlinedFn, EntryFnName);
OffloadInfoManager.registerTargetRegionEntryInfo(
EntryInfo, EntryAddr, OutlinedFnID,
- OffloadEntriesInfoManager::OMPTargetRegionEntryTargetRegion);
+ OffloadEntriesInfoManager::OMPTargetRegionEntryTargetRegion,
+ EntryInfoStorage);
return OutlinedFnID;
}
@@ -5094,7 +5097,8 @@ static void emitTargetOutlinedFunction(
TargetRegionEntryInfo &EntryInfo, Function *&OutlinedFn,
Constant *&OutlinedFnID, SmallVectorImpl<Value *> &Inputs,
OpenMPIRBuilder::TargetBodyGenCallbackTy &CBFunc,
- OpenMPIRBuilder::TargetGenArgAccessorsCallbackTy &ArgAccessorFuncCB) {
+ OpenMPIRBuilder::TargetGenArgAccessorsCallbackTy &ArgAccessorFuncCB,
+ OffloadEntriesInfoManager::OffloadEntryInfoTargetRegion *EntryInfoStorage) {
OpenMPIRBuilder::FunctionGenCallback &&GenerateOutlinedFunction =
[&OMPBuilder, &Builder, &Inputs, &CBFunc,
@@ -5104,7 +5108,8 @@ static void emitTargetOutlinedFunction(
};
OMPBuilder.emitTargetRegionFunction(EntryInfo, GenerateOutlinedFunction, true,
- OutlinedFn, OutlinedFnID);
+ OutlinedFn, OutlinedFnID,
+ EntryInfoStorage);
}
static void emitTargetCall(OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder,
@@ -5165,7 +5170,8 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTarget(
int32_t NumThreads, SmallVectorImpl<Value *> &Args,
GenMapInfoCallbackTy GenMapInfoCB,
OpenMPIRBuilder::TargetBodyGenCallbackTy CBFunc,
- OpenMPIRBuilder::TargetGenArgAccessorsCallbackTy ArgAccessorFuncCB) {
+ OpenMPIRBuilder::TargetGenArgAccessorsCallbackTy ArgAccessorFuncCB,
+ OffloadEntriesInfoManager::OffloadEntryInfoTargetRegion *EntryInfoStorage) {
if (!updateToLocation(Loc))
return InsertPointTy();
@@ -5174,7 +5180,8 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTarget(
Function *OutlinedFn;
Constant *OutlinedFnID;
emitTargetOutlinedFunction(*this, Builder, EntryInfo, OutlinedFn,
- OutlinedFnID, Args, CBFunc, ArgAccessorFuncCB);
+ OutlinedFnID, Args, CBFunc, ArgAccessorFuncCB,
+ EntryInfoStorage);
if (!Config.isTargetDevice())
emitTargetCall(*this, Builder, AllocaIP, OutlinedFn, OutlinedFnID, NumTeams,
NumThreads, Args, GenMapInfoCB);
@@ -6929,7 +6936,8 @@ void OffloadEntriesInfoManager::initializeTargetRegionEntryInfo(
void OffloadEntriesInfoManager::registerTargetRegionEntryInfo(
TargetRegionEntryInfo EntryInfo, Constant *Addr, Constant *ID,
- OMPTargetRegionEntryKind Flags) {
+ OMPTargetRegionEntryKind Flags,
+ OffloadEntryInfoTargetRegion *EntryInfoStorage) {
assert(EntryInfo.Count == 0 && "expected default EntryInfo");
// Update the EntryInfo with the next available count for this location.
@@ -6953,8 +6961,12 @@ void OffloadEntriesInfoManager::registerTargetRegionEntryInfo(
assert(!hasTargetRegionEntryInfo(EntryInfo) &&
"Target region entry already registered!");
OffloadEntryInfoTargetRegion Entry(OffloadingEntriesNum, Addr, ID, Flags);
- OffloadEntriesTargetRegion[EntryInfo] = Entry;
- ++OffloadingEntriesNum;
+ if (EntryInfoStorage) {
+ *EntryInfoStorage = Entry;
+ } else {
+ OffloadEntriesTargetRegion[EntryInfo] = Entry;
+ ++OffloadingEntriesNum;
+ }
}
incrementTargetRegionEntryInfoCount(EntryInfo);
}
>From ac2df60c670a8118a3b4090f0c592ce4838a79a5 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Sun, 14 Jan 2024 01:29:19 +0000
Subject: [PATCH 2/3] [mlir][Target][LLVM] Add offload utility class
This patch adds the `OffloadHandler` utility class for creating LLVM offload
entries.
LLVM offload entries hold information on offload symbols; for example, for a
GPU kernel, this includes its host address to identify the kernel and the kernel
identifier in the binary. Arrays of offload entries can be used to register
functions within the CUDA/HIP runtime. Libomptarget also uses these entries to
register OMP target offload kernels and variables.
This patch is 1/4 on introducing the `OffloadEmbeddingAttr` GPU translation
attribute.
---
mlir/include/mlir/Target/LLVM/Offload.h | 61 ++++++++++++
mlir/lib/Target/LLVM/CMakeLists.txt | 2 +
mlir/lib/Target/LLVM/Offload.cpp | 111 ++++++++++++++++++++++
mlir/unittests/Target/LLVM/CMakeLists.txt | 1 +
mlir/unittests/Target/LLVM/Offload.cpp | 49 ++++++++++
5 files changed, 224 insertions(+)
create mode 100644 mlir/include/mlir/Target/LLVM/Offload.h
create mode 100644 mlir/lib/Target/LLVM/Offload.cpp
create mode 100644 mlir/unittests/Target/LLVM/Offload.cpp
diff --git a/mlir/include/mlir/Target/LLVM/Offload.h b/mlir/include/mlir/Target/LLVM/Offload.h
new file mode 100644
index 00000000000000..7b705667d477d2
--- /dev/null
+++ b/mlir/include/mlir/Target/LLVM/Offload.h
@@ -0,0 +1,61 @@
+//===- Offload.h - LLVM Target Offload --------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares LLVM target offload utility classes.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TARGET_LLVM_OFFLOAD_H
+#define MLIR_TARGET_LLVM_OFFLOAD_H
+
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/StringRef.h"
+
+namespace llvm {
+class Constant;
+class GlobalVariable;
+class Module;
+} // namespace llvm
+
+namespace mlir {
+namespace LLVM {
+/// `OffloadHandler` is a utility class for creating LLVM offload entries. LLVM
+/// offload entries hold information on offload symbols; for example, for a GPU
+/// kernel, this includes its host address to identify the kernel and the kernel
+/// identifier in the binary. Arrays of offload entries can be used to register
+/// functions within the CUDA/HIP runtime. Libomptarget also uses these entries
+/// to register OMP target offload kernels and variables.
+class OffloadHandler {
+public:
+ using OffloadEntryArray =
+ std::pair<llvm::GlobalVariable *, llvm::GlobalVariable *>;
+ OffloadHandler(llvm::Module &module) : module(module) {}
+
+ /// Returns the begin symbol name used in the entry array.
+ static std::string getBeginSymbol(StringRef suffix);
+
+ /// Returns the end symbol name used in the entry array.
+ static std::string getEndSymbol(StringRef suffix);
+
+ /// Returns the entry array if it exists or a pair of null pointers.
+ OffloadEntryArray getEntryArray(StringRef suffix);
+
+ /// Emits an empty array of offloading entries.
+ OffloadEntryArray emitEmptyEntryArray(StringRef suffix);
+
+ /// Inserts an offloading entry into an existing entry array. This method
+ /// returns failure if the entry array hasn't been declared.
+ LogicalResult insertOffloadEntry(StringRef suffix, llvm::Constant *entry);
+
+protected:
+ llvm::Module &module;
+};
+} // namespace LLVM
+} // namespace mlir
+
+#endif // MLIR_TARGET_LLVM_OFFLOAD_H
diff --git a/mlir/lib/Target/LLVM/CMakeLists.txt b/mlir/lib/Target/LLVM/CMakeLists.txt
index cc2c3a00a02eaf..241a6c64dd868f 100644
--- a/mlir/lib/Target/LLVM/CMakeLists.txt
+++ b/mlir/lib/Target/LLVM/CMakeLists.txt
@@ -1,5 +1,6 @@
add_mlir_library(MLIRTargetLLVM
ModuleToObject.cpp
+ Offload.cpp
ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Target/LLVM
@@ -16,6 +17,7 @@ add_mlir_library(MLIRTargetLLVM
Passes
Support
Target
+ FrontendOffloading
LINK_LIBS PUBLIC
MLIRExecutionEngineUtils
MLIRTargetLLVMIRExport
diff --git a/mlir/lib/Target/LLVM/Offload.cpp b/mlir/lib/Target/LLVM/Offload.cpp
new file mode 100644
index 00000000000000..81ba12403bfb99
--- /dev/null
+++ b/mlir/lib/Target/LLVM/Offload.cpp
@@ -0,0 +1,111 @@
+//===- Offload.cpp - LLVM Target Offload ------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines LLVM target offload utility classes.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Target/LLVM/Offload.h"
+#include "llvm/Frontend/Offloading/Utility.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Module.h"
+
+using namespace mlir;
+using namespace mlir::LLVM;
+
+std::string OffloadHandler::getBeginSymbol(StringRef suffix) {
+ return ("__begin_offload_" + suffix).str();
+}
+
+std::string OffloadHandler::getEndSymbol(StringRef suffix) {
+ return ("__end_offload_" + suffix).str();
+}
+
+namespace {
+/// Returns the type of the entry array.
+llvm::ArrayType *getEntryArrayType(llvm::Module &module, size_t numElems) {
+ return llvm::ArrayType::get(llvm::offloading::getEntryTy(module), numElems);
+}
+
+/// Creates the initializer of the entry array.
+llvm::Constant *getEntryArrayBegin(llvm::Module &module,
+ ArrayRef<llvm::Constant *> entries) {
+ // If there are no entries return a constant zero initializer.
+ llvm::ArrayType *arrayTy = getEntryArrayType(module, entries.size());
+ return entries.empty() ? llvm::ConstantAggregateZero::get(arrayTy)
+ : llvm::ConstantArray::get(arrayTy, entries);
+}
+
+/// Computes the end position of the entry array.
+llvm::Constant *getEntryArrayEnd(llvm::Module &module,
+ llvm::GlobalVariable *begin, size_t numElems) {
+ llvm::Type *intTy = module.getDataLayout().getIntPtrType(module.getContext());
+ return llvm::ConstantExpr::getGetElementPtr(
+ llvm::offloading::getEntryTy(module), begin,
+ ArrayRef<llvm::Constant *>({llvm::ConstantInt::get(intTy, numElems)}),
+ true);
+}
+} // namespace
+
+OffloadHandler::OffloadEntryArray
+OffloadHandler::getEntryArray(StringRef suffix) {
+ llvm::GlobalVariable *beginGV =
+ module.getGlobalVariable(getBeginSymbol(suffix), true);
+ llvm::GlobalVariable *endGV =
+ module.getGlobalVariable(getEndSymbol(suffix), true);
+ return {beginGV, endGV};
+}
+
+OffloadHandler::OffloadEntryArray
+OffloadHandler::emitEmptyEntryArray(StringRef suffix) {
+ llvm::ArrayType *arrayTy = getEntryArrayType(module, 0);
+ auto *beginGV = new llvm::GlobalVariable(
+ module, arrayTy, /*isConstant=*/true, llvm::GlobalValue::InternalLinkage,
+ getEntryArrayBegin(module, {}), getBeginSymbol(suffix));
+ auto *endGV = new llvm::GlobalVariable(
+ module, llvm::PointerType::get(module.getContext(), 0),
+ /*isConstant=*/true, llvm::GlobalValue::InternalLinkage,
+ getEntryArrayEnd(module, beginGV, 0), getEndSymbol(suffix));
+ return {beginGV, endGV};
+}
+
+LogicalResult OffloadHandler::insertOffloadEntry(StringRef suffix,
+ llvm::Constant *entry) {
+ // Get the begin and end symbols to the entry array.
+ std::string beginSymId = getBeginSymbol(suffix);
+ llvm::GlobalVariable *beginGV = module.getGlobalVariable(beginSymId, true);
+ llvm::GlobalVariable *endGV =
+ module.getGlobalVariable(getEndSymbol(suffix), true);
+ // Fail if the symbols are missing.
+ if (!beginGV || !endGV)
+ return failure();
+ // Create the entry initializer.
+ assert(beginGV->getInitializer() && "entry array initializer is missing.");
+ // Add existing entries into the new entry array.
+ SmallVector<llvm::Constant *> entries;
+ if (auto beginInit = dyn_cast_or_null<llvm::ConstantAggregate>(
+ beginGV->getInitializer())) {
+ for (unsigned i = 0; i < beginInit->getNumOperands(); ++i)
+ entries.push_back(beginInit->getOperand(i));
+ }
+ // Add the new entry.
+ entries.push_back(entry);
+ // Create a global holding the new updated set of entries.
+ auto *arrayTy = llvm::ArrayType::get(llvm::offloading::getEntryTy(module),
+ entries.size());
+ auto *entryArr = new llvm::GlobalVariable(
+ module, arrayTy, /*isConstant=*/true, llvm::GlobalValue::InternalLinkage,
+ getEntryArrayBegin(module, entries), beginSymId, endGV);
+ // Replace the old entry array variable withe new one.
+ beginGV->replaceAllUsesWith(entryArr);
+ beginGV->eraseFromParent();
+ entryArr->setName(beginSymId);
+ // Update the end symbol.
+ endGV->setInitializer(getEntryArrayEnd(module, entryArr, entries.size()));
+ return success();
+}
diff --git a/mlir/unittests/Target/LLVM/CMakeLists.txt b/mlir/unittests/Target/LLVM/CMakeLists.txt
index 6d612548a94c0f..d04f38ddddfacf 100644
--- a/mlir/unittests/Target/LLVM/CMakeLists.txt
+++ b/mlir/unittests/Target/LLVM/CMakeLists.txt
@@ -1,4 +1,5 @@
add_mlir_unittest(MLIRTargetLLVMTests
+ Offload.cpp
SerializeNVVMTarget.cpp
SerializeROCDLTarget.cpp
SerializeToLLVMBitcode.cpp
diff --git a/mlir/unittests/Target/LLVM/Offload.cpp b/mlir/unittests/Target/LLVM/Offload.cpp
new file mode 100644
index 00000000000000..375edc2e9614d3
--- /dev/null
+++ b/mlir/unittests/Target/LLVM/Offload.cpp
@@ -0,0 +1,49 @@
+//===- Offload.cpp ----------------------------------------------*- C++ -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Target/LLVM/Offload.h"
+#include "llvm/Frontend/Offloading/Utility.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Module.h"
+
+#include "gmock/gmock.h"
+
+using namespace llvm;
+
+TEST(MLIRTarget, OffloadAPI) {
+ using OffloadEntryArray = mlir::LLVM::OffloadHandler::OffloadEntryArray;
+ LLVMContext llvmContext;
+ Module llvmModule("offload", llvmContext);
+ mlir::LLVM::OffloadHandler handler(llvmModule);
+ StringRef suffix = ".mlir";
+ // Check there's no entry array with `.mlir` suffix.
+ OffloadEntryArray entryArray = handler.getEntryArray(suffix);
+ EXPECT_EQ(entryArray, OffloadEntryArray());
+ // Emit the entry array.
+ handler.emitEmptyEntryArray(suffix);
+ // Check there's an entry array with `.mlir` suffix.
+ entryArray = handler.getEntryArray(suffix);
+ ASSERT_NE(entryArray.first, nullptr);
+ ASSERT_NE(entryArray.second, nullptr);
+ // Check the array contains no entries.
+ auto *zeroInitializer = dyn_cast_or_null<ConstantAggregateZero>(
+ entryArray.first->getInitializer());
+ ASSERT_NE(zeroInitializer, nullptr);
+ // Insert an empty entries.
+ auto emptyEntry =
+ ConstantAggregateZero::get(offloading::getEntryTy(llvmModule));
+ ASSERT_TRUE(succeeded(handler.insertOffloadEntry(suffix, emptyEntry)));
+ // Check there's an entry in the entry array with `.mlir` suffix.
+ entryArray = handler.getEntryArray(suffix);
+ ASSERT_NE(entryArray.first, nullptr);
+ Constant *arrayInitializer = entryArray.first->getInitializer();
+ ASSERT_NE(arrayInitializer, nullptr);
+ auto *arrayTy = dyn_cast_or_null<ArrayType>(arrayInitializer->getType());
+ ASSERT_NE(arrayTy, nullptr);
+ EXPECT_EQ(arrayTy->getNumElements(), 1u);
+}
>From 5c08090f05704ae1464fa2b09ce21fb4ddc9471d Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Mon, 15 Jan 2024 16:44:27 +0000
Subject: [PATCH 3/3] [mlir][OpenMP] Add the `tgt_entry_info` attribute
This patch adds the `omp.tgt_entry_info` attribute. This attribute provides
information to identify offload entries uniquely and partially reflects the
information in the `llvm::TargetRegionEntryInfo` struct.
An `info` parameter was added to `TargetOp` to specify the offload entry
information from the operation explicitly. Both the host and device versions
must have the same `info` attribute; otherwise, the constructs won't correctly
map between each other.
This patch is required to enable JIT compilation for the OMP dialect, as
the entry array has to be fully constructed in the IR instead of using
sections to implicitly construct it.
Note: Ignore the base commits.
---
flang/lib/Lower/OpenMP.cpp | 2 +-
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 34 ++++++++++-
.../LLVMIR/Dialect/OpenMP/CMakeLists.txt | 1 +
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 61 ++++++++++++++++++-
.../Target/LLVMIR/omptarget-entry-info.mlir | 48 +++++++++++++++
5 files changed, 142 insertions(+), 4 deletions(-)
create mode 100644 mlir/test/Target/LLVMIR/omptarget-entry-info.mlir
diff --git a/flang/lib/Lower/OpenMP.cpp b/flang/lib/Lower/OpenMP.cpp
index 4f7c99a6d2b840..2944187aaef7aa 100644
--- a/flang/lib/Lower/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP.cpp
@@ -2802,7 +2802,7 @@ genTargetOp(Fortran::lower::AbstractConverter &converter,
auto targetOp = converter.getFirOpBuilder().create<mlir::omp::TargetOp>(
currentLocation, ifClauseOperand, deviceOperand, threadLimitOperand,
- nowaitAttr, mapOperands);
+ nowaitAttr, mapOperands, mlir::omp::TargetRegionEntryInfoAttr());
genBodyOfTargetOp(converter, eval, targetOp, mapSymTypes, mapSymLocs,
mapSymbols, currentLocation);
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index d614f2666a85ab..89fea898497be4 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -81,6 +81,33 @@ def TargetAttr : OpenMP_Attr<"Target", "target"> {
let assemblyFormat = "`<` struct(params) `>`";
}
+def TargetRegionEntryInfoAttr :
+ OpenMP_Attr<"TargetRegionEntryInfo", "tgt_entry_info"> {
+ let description = [{
+ The `tgt_entry_info` attribute provides information to identify offload
+ entries uniquely. This attribute partially reflects the information in the
+ `llvm::TargetRegionEntryInfo` struct.
+
+ The information in this attribute can be used to generate the unique
+ identifier used to refer the offload symbol.
+
+ The optional `section` parameter can be used to emit the entry in an specific
+ entry rather than using the `omp_offloading_entries` data section. This
+ array has to be created before translating the Op generating the entry.
+
+ Example:
+ ```mlir
+ omp.tgt_entry_info<deviceID = 1, fileID = 0, line = 1, section = @entryArray>
+ ```
+ }];
+ let parameters = (ins
+ "unsigned":$deviceID,
+ "unsigned":$fileID,
+ "unsigned":$line,
+ OptionalParameter<"FlatSymbolRefAttr">:$section
+ );
+ let assemblyFormat = "`<` struct(params) `>`";
+}
class OpenMP_Op<string mnemonic, list<Trait> traits = []> :
Op<OpenMP_Dialect, mnemonic, traits>;
@@ -1443,6 +1470,9 @@ def TargetOp : OpenMP_Op<"target",[IsolatedFromAbove, OutlineableOpenMPOpInterfa
The optional $nowait elliminates the implicit barrier so the parent task can make progress
even if the target task is not yet completed.
+ The optional $info paremeter specifies the information that should be used to
+ create the offload entry in the IR.
+
TODO: is_device_ptr, depend, defaultmap, in_reduction
}];
@@ -1451,7 +1481,8 @@ def TargetOp : OpenMP_Op<"target",[IsolatedFromAbove, OutlineableOpenMPOpInterfa
Optional<AnyInteger>:$device,
Optional<AnyInteger>:$thread_limit,
UnitAttr:$nowait,
- Variadic<AnyType>:$map_operands);
+ Variadic<AnyType>:$map_operands,
+ OptionalAttr<TargetRegionEntryInfoAttr>:$targetRegionEntryInfo);
let regions = (region AnyRegion:$region);
@@ -1461,6 +1492,7 @@ def TargetOp : OpenMP_Op<"target",[IsolatedFromAbove, OutlineableOpenMPOpInterfa
| `thread_limit` `(` $thread_limit `:` type($thread_limit) `)`
| `nowait` $nowait
| `map_entries` `(` custom<MapEntries>($map_operands, type($map_operands)) `)`
+ | `info` `=` qualified($targetRegionEntryInfo)
) $region attr-dict
}];
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/CMakeLists.txt b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/CMakeLists.txt
index 0a5d7c6e22058d..744ec7b1ae3cf4 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/CMakeLists.txt
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/CMakeLists.txt
@@ -9,6 +9,7 @@ add_mlir_translation_library(MLIROpenMPToLLVMIRTranslation
MLIRLLVMDialect
MLIROpenMPDialect
MLIRSupport
+ MLIRTargetLLVM
MLIRTargetLLVMIRExport
MLIRTransformUtils
)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index e7aebc3ce4be56..5492d828a99a3f 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -18,12 +18,14 @@
#include "mlir/IR/Operation.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Support/LogicalResult.h"
+#include "mlir/Target/LLVM/Offload.h"
#include "mlir/Target/LLVMIR/Dialect/OpenMPCommon.h"
#include "mlir/Target/LLVMIR/ModuleTranslation.h"
#include "mlir/Transforms/RegionUtils.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Frontend/Offloading/Utility.h"
#include "llvm/Frontend/OpenMP/OMPConstants.h"
#include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
#include "llvm/IR/DebugInfoMetadata.h"
@@ -2065,6 +2067,25 @@ LogicalResult convertFlagsAttr(Operation *op, mlir::omp::FlagsAttr attribute,
return success();
}
+static bool
+getAndRegisterTargetEntryUniqueInfo(llvm::TargetRegionEntryInfo &targetInfo,
+ LLVM::ModuleTranslation &moduleTranslation,
+ omp::TargetOp targetOp,
+ llvm::StringRef parentName = "") {
+ omp::TargetRegionEntryInfoAttr infoAttr =
+ targetOp.getTargetRegionEntryInfoAttr();
+ if (!infoAttr)
+ return false;
+ targetInfo =
+ llvm::TargetRegionEntryInfo(parentName, infoAttr.getDeviceID(),
+ infoAttr.getFileID(), infoAttr.getLine());
+ llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
+ if (ompBuilder->Config.isTargetDevice() && ompBuilder->Config.isGPU())
+ ompBuilder->OffloadInfoManager.initializeTargetRegionEntryInfo(targetInfo,
+ 0);
+ return true;
+}
+
static bool getTargetEntryUniqueInfo(llvm::TargetRegionEntryInfo &targetInfo,
omp::TargetOp targetOp,
llvm::StringRef parentName = "") {
@@ -2371,7 +2392,9 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
llvm::TargetRegionEntryInfo entryInfo;
- if (!getTargetEntryUniqueInfo(entryInfo, targetOp, parentName))
+ if (!(getAndRegisterTargetEntryUniqueInfo(entryInfo, moduleTranslation,
+ targetOp, parentName) ||
+ getTargetEntryUniqueInfo(entryInfo, targetOp, parentName)))
return failure();
int32_t defaultValTeams = -1;
@@ -2434,15 +2457,49 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
kernelInput.push_back(mapData.OriginalValue[i]);
}
+ llvm::OffloadEntriesInfoManager::OffloadEntryInfoTargetRegion entryRegionInfo;
+ // Determine whether the entry is going to be handled by
+ // `OffloadEntriesInfoManager` or by this method. If `entryArraySection` is
+ // null then it's handled by `OffloadEntriesInfoManager`
+ omp::TargetRegionEntryInfoAttr regionInfoAttr =
+ targetOp.getTargetRegionEntryInfoAttr();
+ FlatSymbolRefAttr entryArraySection =
+ regionInfoAttr ? regionInfoAttr.getSection() : FlatSymbolRefAttr();
+
+ // Create the target region
builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createTarget(
ompLoc, allocaIP, builder.saveIP(), entryInfo, defaultValTeams,
- defaultValThreads, kernelInput, genMapInfoCB, bodyCB, argAccessorCB));
+ defaultValThreads, kernelInput, genMapInfoCB, bodyCB, argAccessorCB,
+ entryArraySection ? &entryRegionInfo : nullptr));
// Remap access operations to declare target reference pointers for the
// device, essentially generating extra loadop's as necessary
if (moduleTranslation.getOpenMPBuilder()->Config.isTargetDevice())
handleDeclareTargetMapVar(mapData, moduleTranslation, builder);
+ // Return early if the target op it's being emitted for a device or if the
+ // entry is handled by `OffloadEntriesInfoManager`
+ llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
+ if (ompBuilder->Config.isTargetDevice() || !entryArraySection)
+ return bodyGenStatus;
+
+ assert(entryRegionInfo.isValid() && "invalid target entry region");
+
+ auto regionAddrGV =
+ dyn_cast_or_null<llvm::GlobalValue>(entryRegionInfo.getAddress());
+ assert(regionAddrGV && "missing reggion address");
+
+ // Emit the offload entry.
+ llvm::Module &llvmModule = *moduleTranslation.getLLVMModule();
+ LLVM::OffloadHandler offloadHandler(llvmModule);
+ std::pair<llvm::Constant *, llvm::GlobalVariable *> entryInit =
+ llvm::offloading::getOffloadingEntryInitializer(
+ llvmModule, entryRegionInfo.getID(), regionAddrGV->getName(), 0,
+ entryRegionInfo.getFlags(), 0);
+
+ if (failed(offloadHandler.insertOffloadEntry(entryArraySection.getValue(),
+ entryInit.first)))
+ targetOp.emitError("failed to insert the entry");
return bodyGenStatus;
}
diff --git a/mlir/test/Target/LLVMIR/omptarget-entry-info.mlir b/mlir/test/Target/LLVMIR/omptarget-entry-info.mlir
new file mode 100644
index 00000000000000..dafae687022291
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-entry-info.mlir
@@ -0,0 +1,48 @@
+// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck %s
+
+!EntryArray = !llvm.struct<(!llvm.ptr, !llvm.ptr, i64, i32, i32)>
+// CHECK: @__begin_offload_omp = internal constant [2 x %{{.*}}] [%{{.*}} { ptr @[[TGT_OP_1:.*]], ptr @[[TGT_OP_1_NAME:.*]], i64 0, i32 0, i32 0 }, %{{.*}} { ptr @[[TGT_OP_2:.*]], ptr @[[TGT_OP_2_NAME:.*]], i64 0, i32 0, i32 0 }]
+// CHECK: @__end_offload_omp = constant ptr getelementptr inbounds (%{{.*}}, ptr @__begin_offload_omp, i64 2)
+// CHECK: @[[TGT_OP_1]] = weak constant i8 0
+// CHECK: @[[TGT_OP_1_NAME]] = internal unnamed_addr constant [{{.*}} x i8] c"{{.*}}0_0_main_l0\00"
+// CHECK: @[[TGT_OP_2]] = weak constant i8 0
+// CHECK: @[[TGT_OP_2_NAME]] = internal unnamed_addr constant [{{.*}} x i8] c"{{.*}}0_0_main_l1\00"
+// CHECK: define void @main() {
+// CHECK: %{{.*}} = call i32 @__tgt_target_kernel(ptr @{{.*}}, i64 -1, i32 -1, i32 0, ptr @[[TGT_OP_1]], ptr %{{.*}})
+// CHECK: %{{.*}} = call i32 @__tgt_target_kernel(ptr @{{.*}}, i64 -1, i32 -1, i32 0, ptr @[[TGT_OP_2]], ptr %{{.*}})
+// CHECK: }
+// CHECK-LABEL: define internal void @{{.*}}0_0_main_l0() {
+// CHECK-LABEL: define internal void @{{.*}}0_0_main_l1() {
+module attributes {omp.is_target_device = false, omp.is_gpu = false} {
+ llvm.mlir.global constant @__begin_offload_omp() : !llvm.array<0 x !EntryArray> {
+ %zero = llvm.mlir.zero : !llvm.array<0 x !EntryArray>
+ llvm.return %zero : !llvm.array<0 x !EntryArray>
+ }
+ llvm.mlir.global constant @__end_offload_omp() : !llvm.ptr {
+ %array = llvm.mlir.addressof @__begin_offload_omp : !llvm.ptr
+ llvm.return %array : !llvm.ptr
+ }
+ llvm.func @main() {
+ omp.target info = #omp.tgt_entry_info<deviceID = 0, fileID = 0, line = 0, section = @omp> {
+ omp.terminator
+ }
+ omp.target info = #omp.tgt_entry_info<deviceID = 0, fileID = 0, line = 1, section = @omp> {
+ omp.terminator
+ }
+ llvm.return
+ }
+}
+
+// -----
+
+// CHECK: @[[TGT_OP:.*]] = weak constant i8 0
+// CHECK: @[[TGT_OP_NAME:.*]] = internal unnamed_addr constant [{{.*}} x i8] c"{{.*}}0_0_main_l0\00"
+// CHECK: @{{.*}} = weak constant %{{.*}} { ptr @[[TGT_OP]], ptr @[[TGT_OP_NAME]], i64 0, i32 0, i32 0 }, section "omp_offloading_entries", align 1
+module attributes {omp.is_target_device = false, omp.is_gpu = false} {
+ llvm.func @main() {
+ omp.target info = #omp.tgt_entry_info<deviceID = 0, fileID = 0, line = 0> {
+ omp.terminator
+ }
+ llvm.return
+ }
+}
More information about the flang-commits
mailing list