Skip to content

[Offload] Properly guard modifications to the RPC device array #126790

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Feb 11, 2025
Merged

[Offload] Properly guard modifications to the RPC device array #126790

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 9 additions & 3 deletions offload/plugins-nextgen/common/include/RPC.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,9 @@ struct RPCServerTy {
/// Array of associated devices. These must be alive as long as the server is.
std::unique_ptr<plugin::GenericDeviceTy *[]> Devices;

/// Mutex that guards accesses to the buffers and device array.
std::mutex BufferMutex{};

/// A helper class for running the user thread that handles the RPC interface.
/// Because we only need to check the RPC server while any kernels are
/// working, we track submission / completion events to allow the thread to
Expand All @@ -90,6 +93,9 @@ struct RPCServerTy {
std::condition_variable CV;
std::mutex Mutex;

/// A reference to the main server's mutex.
std::mutex &BufferMutex;

/// A reference to all the RPC interfaces that the server is handling.
llvm::ArrayRef<void *> Buffers;

Expand All @@ -98,9 +104,9 @@ struct RPCServerTy {

/// Initialize the worker thread to run in the background.
ServerThread(void *Buffers[], plugin::GenericDeviceTy *Devices[],
size_t Length)
: Running(false), NumUsers(0), CV(), Mutex(), Buffers(Buffers, Length),
Devices(Devices, Length) {}
size_t Length, std::mutex &BufferMutex)
: Running(false), NumUsers(0), CV(), Mutex(), BufferMutex(BufferMutex),
Buffers(Buffers, Length), Devices(Devices, Length) {}

~ServerThread() { assert(!Running && "Thread not shut down explicitly\n"); }

Expand Down
5 changes: 4 additions & 1 deletion offload/plugins-nextgen/common/src/RPC.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,7 @@ void RPCServerTy::ServerThread::run() {
Lock.unlock();
while (NumUsers.load(std::memory_order_relaxed) > 0 &&
Running.load(std::memory_order_relaxed)) {
std::lock_guard<decltype(Mutex)> Lock(BufferMutex);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How often do we hit this line?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Every time we do a full scan of all the ports, so very often but I don't think there's a way to avoid that overhead.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So that is when we have something on the device actually using the RPC mechanism then we will hit this overhead, right?

Copy link
Contributor Author

@jhuber6 jhuber6 Feb 11, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, it might make it slightly slower if you're hitting RPC hard. As in, every thread constantly submitting requests.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok, but if this only occurs when you actually hit RPC I'm much less concerned and think this is fine for now.

for (const auto &[Buffer, Device] : llvm::zip_equal(Buffers, Devices)) {
if (!Buffer || !Device)
continue;
Expand All @@ -146,7 +147,7 @@ RPCServerTy::RPCServerTy(plugin::GenericPluginTy &Plugin)
Devices(std::make_unique<plugin::GenericDeviceTy *[]>(
Plugin.getNumDevices())),
Thread(new ServerThread(Buffers.get(), Devices.get(),
Plugin.getNumDevices())) {}
Plugin.getNumDevices(), BufferMutex)) {}

llvm::Error RPCServerTy::startThread() {
Thread->startThread();
Expand Down Expand Up @@ -187,13 +188,15 @@ Error RPCServerTy::initDevice(plugin::GenericDeviceTy &Device,
if (auto Err = Device.dataSubmit(ClientGlobal.getPtr(), &client,
sizeof(rpc::Client), nullptr))
return Err;
std::lock_guard<decltype(BufferMutex)> Lock(BufferMutex);
Buffers[Device.getDeviceId()] = RPCBuffer;
Devices[Device.getDeviceId()] = &Device;

return Error::success();
}

Error RPCServerTy::deinitDevice(plugin::GenericDeviceTy &Device) {
std::lock_guard<decltype(BufferMutex)> Lock(BufferMutex);
Device.free(Buffers[Device.getDeviceId()], TARGET_ALLOC_HOST);
Buffers[Device.getDeviceId()] = nullptr;
Devices[Device.getDeviceId()] = nullptr;
Expand Down