[llvm-branch-commits] [llvm] [offload][omp] Route RPC callback registration through liboffload (PR #221840)

[Alex Duran via llvm-branch-commits]

https://github.com/adurang updated https://github.com/llvm/llvm-project/pull/221840

>From 78656581e687f854606f7f822b4d912e3131db4f Mon Sep 17 00:00:00 2001
From: "Duran, Alex" <alejandro.duran at intel.com>
Date: Mon, 7 Sep 2026 14:27:10 -0700
Subject: [PATCH] [offload][omp] Route RPC callback registration through
 liboffload

Replace __tgt_register_rpc_callback's direct iteration over plugins
with olIteratePlatforms + olPlatformRegisterRPCCallback, and drop the
now-unused RPCServerTy::registerCallback export. Move the
initialized/has-devices guard that used to live in libomptarget into
olPlatformRegisterRPCCallback_impl.
---
 offload/liboffload/exports             |  1 -
 offload/liboffload/src/OffloadImpl.cpp |  4 ++++
 offload/libomptarget/interface.cpp     | 10 +++++++---
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/offload/liboffload/exports b/offload/liboffload/exports
index d39461bbe6ac3..f7f85e452af80 100644
--- a/offload/liboffload/exports
+++ b/offload/liboffload/exports
@@ -7,7 +7,6 @@ global:
   extern "C++" {
     error::OffloadError::ID;
     "error::OffloadErrCategory()";
-    "llvm::omp::target::RPCServerTy::registerCallback(unsigned int (*)(void*, unsigned int))";
     "llvm::omp::target::plugin::GenericPluginTy::async_barrier(omp_interop_val_t*)";
     "llvm::omp::target::plugin::GenericPluginTy::create_interop(int, int, interop_spec_t*)";
     "llvm::omp::target::plugin::GenericPluginTy::data_alloc(int, long, void*, int)";
diff --git a/offload/liboffload/src/OffloadImpl.cpp b/offload/liboffload/src/OffloadImpl.cpp
index 352944852d214..16e1c40d273e8 100644
--- a/offload/liboffload/src/OffloadImpl.cpp
+++ b/offload/liboffload/src/OffloadImpl.cpp
@@ -465,6 +465,10 @@ Error olGetPlatformInfoSize_impl(ol_platform_handle_t Platform,
 
 Error olPlatformRegisterRPCCallback_impl(ol_platform_handle_t Platform,
                                          ol_platform_rpc_cb_t Callback) {
+  if (!Platform->Plugin || !Platform->Plugin->is_initialized() ||
+      Platform->Plugin->getNumDevices() == 0)
+    return Error::success();
+
   Platform->Plugin->getRPCServer().registerCallback(Callback);
   return Error::success();
 }
diff --git a/offload/libomptarget/interface.cpp b/offload/libomptarget/interface.cpp
index 5b5cc310b1b31..1afc2a7fe51eb 100644
--- a/offload/libomptarget/interface.cpp
+++ b/offload/libomptarget/interface.cpp
@@ -652,7 +652,11 @@ EXTERN void __tgt_register_rpc_callback(unsigned (*Callback)(void *,
   if (!PM)
     return;
 
-  for (auto &Plugin : PM->plugins())
-    if (Plugin.is_initialized() && Plugin.getNumDevices() > 0)
-      Plugin.getRPCServer().registerCallback(Callback);
+  olIteratePlatforms(
+      [](ol_platform_handle_t Platform, void *Data) {
+        olPlatformRegisterRPCCallback(
+            Platform, reinterpret_cast<ol_platform_rpc_cb_t>(Data));
+        return true;
+      },
+      reinterpret_cast<void *>(Callback));
 }