Skip to content

Commit 6f7ca4c

Browse files
[Offload] Add MPI Proxy Plugin
Co-authored-by: Guilherme Valarini <[email protected]>
1 parent 6ca59aa commit 6f7ca4c

34 files changed

+4756
-263
lines changed

offload/CMakeLists.txt

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,7 @@ if(DEFINED LIBOMPTARGET_BUILD_CUDA_PLUGIN OR
145145
message(WARNING "Option removed, use 'LIBOMPTARGET_PLUGINS_TO_BUILD' instead")
146146
endif()
147147

148-
set(LIBOMPTARGET_ALL_PLUGIN_TARGETS amdgpu cuda host)
148+
set(LIBOMPTARGET_ALL_PLUGIN_TARGETS mpi amdgpu cuda host)
149149
set(LIBOMPTARGET_PLUGINS_TO_BUILD "all" CACHE STRING
150150
"Semicolon-separated list of plugins to use: cuda, amdgpu, host or \"all\".")
151151

@@ -200,8 +200,10 @@ set (LIBOMPTARGET_ALL_TARGETS "${LIBOMPTARGET_ALL_TARGETS} powerpc64-ibm-linux-g
200200
set (LIBOMPTARGET_ALL_TARGETS "${LIBOMPTARGET_ALL_TARGETS} powerpc64-ibm-linux-gnu-LTO")
201201
set (LIBOMPTARGET_ALL_TARGETS "${LIBOMPTARGET_ALL_TARGETS} x86_64-unknown-linux-gnu")
202202
set (LIBOMPTARGET_ALL_TARGETS "${LIBOMPTARGET_ALL_TARGETS} x86_64-unknown-linux-gnu-LTO")
203+
set (LIBOMPTARGET_ALL_TARGETS "${LIBOMPTARGET_ALL_TARGETS} x86_64-unknown-linux-gnu-mpi")
203204
set (LIBOMPTARGET_ALL_TARGETS "${LIBOMPTARGET_ALL_TARGETS} nvptx64-nvidia-cuda")
204205
set (LIBOMPTARGET_ALL_TARGETS "${LIBOMPTARGET_ALL_TARGETS} nvptx64-nvidia-cuda-LTO")
206+
set (LIBOMPTARGET_ALL_TARGETS "${LIBOMPTARGET_ALL_TARGETS} nvptx64-nvidia-cuda-mpi")
205207
set (LIBOMPTARGET_ALL_TARGETS "${LIBOMPTARGET_ALL_TARGETS} nvptx64-nvidia-cuda-JIT-LTO")
206208
set (LIBOMPTARGET_ALL_TARGETS "${LIBOMPTARGET_ALL_TARGETS} s390x-ibm-linux-gnu")
207209
set (LIBOMPTARGET_ALL_TARGETS "${LIBOMPTARGET_ALL_TARGETS} s390x-ibm-linux-gnu-LTO")
@@ -363,6 +365,8 @@ set(LIBOMPTARGET_LLVM_LIBRARY_DIR "${LLVM_LIBRARY_DIR}" CACHE STRING
363365
set(LIBOMPTARGET_LLVM_LIBRARY_INTDIR "${LIBOMPTARGET_INTDIR}" CACHE STRING
364366
"Path to folder where intermediate libraries will be output")
365367

368+
set(LIBOMPTARGET_SRC_DIR ${CMAKE_CURRENT_SOURCE_DIR}/libomptarget)
369+
366370
add_subdirectory(tools/offload-tblgen)
367371

368372
# Build offloading plugins and device RTLs if they are available.

offload/include/omptarget.h

Lines changed: 74 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,77 @@ inline KernelArgsTy CTorDTorKernelArgs = {1, 0, nullptr, nullptr,
109109
nullptr, nullptr, nullptr, nullptr,
110110
0, {0,0,0}, {1, 0, 0}, {1, 0, 0}, 0};
111111

112+
using llvm::SmallVector;
112113
struct DeviceTy;
114+
class AsyncInfoTy;
115+
116+
/// A class manages private arguments in a target region.
117+
class PrivateArgumentManagerTy {
118+
/// A data structure for the information of first-private arguments. We can
119+
/// use this information to optimize data transfer by packing all
120+
/// first-private arguments and transfer them all at once.
121+
struct FirstPrivateArgInfoTy {
122+
/// Host pointer begin
123+
char *HstPtrBegin;
124+
/// Host pointer end
125+
char *HstPtrEnd;
126+
/// The index of the element in \p TgtArgs corresponding to the argument
127+
int Index;
128+
/// Alignment of the entry (base of the entry, not after the entry).
129+
uint32_t Alignment;
130+
/// Size (without alignment, see padding)
131+
uint32_t Size;
132+
/// Padding used to align this argument entry, if necessary.
133+
uint32_t Padding;
134+
/// Host pointer name
135+
map_var_info_t HstPtrName = nullptr;
136+
137+
FirstPrivateArgInfoTy(int Index, void *HstPtr, uint32_t Size,
138+
uint32_t Alignment, uint32_t Padding,
139+
map_var_info_t HstPtrName = nullptr)
140+
: HstPtrBegin(reinterpret_cast<char *>(HstPtr)),
141+
HstPtrEnd(HstPtrBegin + Size), Index(Index), Alignment(Alignment),
142+
Size(Size), Padding(Padding), HstPtrName(HstPtrName) {}
143+
};
144+
145+
/// A vector of target pointers for all private arguments
146+
SmallVector<void *> TgtPtrs;
147+
148+
/// A vector of information of all first-private arguments to be packed
149+
SmallVector<FirstPrivateArgInfoTy> FirstPrivateArgInfo;
150+
/// Host buffer for all arguments to be packed
151+
SmallVector<char> FirstPrivateArgBuffer;
152+
/// The total size of all arguments to be packed
153+
int64_t FirstPrivateArgSize = 0;
154+
155+
/// A reference to the \p DeviceTy object
156+
DeviceTy &Device;
157+
/// A pointer to a \p AsyncInfoTy object
158+
AsyncInfoTy &AsyncInfo;
159+
160+
// TODO: What would be the best value here? Should we make it configurable?
161+
// If the size is larger than this threshold, we will allocate and transfer it
162+
// immediately instead of packing it.
163+
static constexpr const int64_t FirstPrivateArgSizeThreshold = 1024;
164+
165+
public:
166+
/// Constructor
167+
PrivateArgumentManagerTy(DeviceTy &Dev, AsyncInfoTy &AsyncInfo)
168+
: Device(Dev), AsyncInfo(AsyncInfo) {}
169+
170+
/// Add a private argument
171+
int addArg(void *HstPtr, int64_t ArgSize, int64_t ArgOffset,
172+
bool IsFirstPrivate, void *&TgtPtr, int TgtArgsIndex,
173+
map_var_info_t HstPtrName = nullptr,
174+
const bool AllocImmediately = false);
175+
176+
/// Pack first-private arguments, replace place holder pointers in \p TgtArgs,
177+
/// and start the transfer.
178+
int packAndTransfer(SmallVector<void *> &TgtArgs);
179+
180+
/// Free all target memory allocated for private arguments
181+
int free();
182+
};
113183

114184
/// The libomptarget wrapper around a __tgt_async_info object directly
115185
/// associated with a libomptarget layer device. RAII semantics to avoid
@@ -136,8 +206,11 @@ class AsyncInfoTy {
136206
/// Synchronization method to be used.
137207
SyncTy SyncType;
138208

209+
PrivateArgumentManagerTy PrivateArgumentManager;
210+
139211
AsyncInfoTy(DeviceTy &Device, SyncTy SyncType = SyncTy::BLOCKING)
140-
: Device(Device), SyncType(SyncType) {}
212+
: Device(Device), SyncType(SyncType),
213+
PrivateArgumentManager(Device, *this) {}
141214
~AsyncInfoTy() { synchronize(); }
142215

143216
/// Implicit conversion to the __tgt_async_info which is used in the

offload/libomptarget/PluginManager.cpp

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
#include "PluginManager.h"
1414
#include "OffloadPolicy.h"
1515
#include "Shared/Debug.h"
16+
#include "Shared/EnvironmentVar.h"
1617
#include "Shared/Profile.h"
1718
#include "device.h"
1819

@@ -71,6 +72,12 @@ bool PluginManager::initializePlugin(GenericPluginTy &Plugin) {
7172
if (Plugin.is_initialized())
7273
return true;
7374

75+
// Disable Host Plugin when it is needed
76+
IntEnvar DisableHostPlugin("OMPTARGET_DISABLE_HOST_PLUGIN", 0);
77+
if (DisableHostPlugin.get() && !strcmp(Plugin.getName(), "x86_64")) {
78+
return false;
79+
}
80+
7481
if (auto Err = Plugin.init()) {
7582
[[maybe_unused]] std::string InfoMsg = toString(std::move(Err));
7683
DP("Failed to init plugin: %s\n", InfoMsg.c_str());

0 commit comments

Comments
 (0)