https://github.com/adurang updated https://github.com/llvm/llvm-project/pull/226441
>From 7afc5ab56d454beac0393ded9271bfe22c072b77 Mon Sep 17 00:00:00 2001 From: "Duran, Alex" <[email protected]> Date: Mon, 7 Sep 2026 03:17:04 -0700 Subject: [PATCH] [OFFLOAD] Add olIterateCompatibleDevices API --- offload/liboffload/API/Program.td | 15 ++++ offload/liboffload/src/OffloadImpl.cpp | 27 ++++++- offload/unittests/OffloadAPI/CMakeLists.txt | 1 + .../program/olIterateCompatibleDevices.cpp | 73 +++++++++++++++++++ 4 files changed, 114 insertions(+), 2 deletions(-) create mode 100644 offload/unittests/OffloadAPI/program/olIterateCompatibleDevices.cpp diff --git a/offload/liboffload/API/Program.td b/offload/liboffload/API/Program.td index ecc8fb73dad3f..9f945872aec3c 100644 --- a/offload/liboffload/API/Program.td +++ b/offload/liboffload/API/Program.td @@ -46,6 +46,21 @@ def olIsValidBinary : Function { let returns = []; } +def olIterateCompatibleDevices : Function { + let desc = "Iterates over all available devices that are compatible with the binary image pointed to by `ProgData`, calling the callback for each device."; + let details = [ + "The provided `ProgData` will not be loaded onto any device", + "If the user-provided callback returns `false`, the iteration is stopped." + ]; + let params = [ + Param<"const void*", "ProgData", "pointer to the program binary data", PARAM_IN>, + Param<"size_t", "ProgDataSize", "size of the program binary in bytes", PARAM_IN>, + Param<"ol_device_iterate_cb_t", "Callback", "User-provided function called for each compatible device", PARAM_IN>, + Param<"void*", "UserData", "Optional user data to pass to the callback", PARAM_IN_OPTIONAL> + ]; + let returns = []; +} + def olDestroyProgram : Function { let desc = "Destroy the program and free all underlying resources."; let details = []; diff --git a/offload/liboffload/src/OffloadImpl.cpp b/offload/liboffload/src/OffloadImpl.cpp index 17ba5c0b4cef4..cc5e8a2ad304d 100644 --- a/offload/liboffload/src/OffloadImpl.cpp +++ b/offload/liboffload/src/OffloadImpl.cpp @@ -673,9 +673,32 @@ Error olIterateDevices_impl(ol_device_iterate_cb_t Callback, void *UserData) { if (!DevicesOrErr) return DevicesOrErr.takeError(); for (auto &Device : *DevicesOrErr) { - if (!Callback(Device.get(), UserData)) { + if (!Callback(Device.get(), UserData)) + return Error::success(); + } + } + + return Error::success(); +} + +Error olIterateCompatibleDevices_impl(const void *ProgData, size_t ProgDataSize, + ol_device_iterate_cb_t Callback, + void *UserData) { + StringRef Buffer(reinterpret_cast<const char *>(ProgData), ProgDataSize); + + for (auto &Platform : OffloadContext::get().Platforms) { + if (!Platform->Plugin || !Platform->Plugin->isPluginCompatible(Buffer)) + continue; + auto DevicesOrErr = Platform->getDevices(); + if (!DevicesOrErr) + return DevicesOrErr.takeError(); + for (auto &Device : *DevicesOrErr) { + if (!Device->Platform.Plugin->isDeviceCompatible(Device->DeviceNum, + Buffer)) + continue; + + if (!Callback(Device.get(), UserData)) return Error::success(); - } } } diff --git a/offload/unittests/OffloadAPI/CMakeLists.txt b/offload/unittests/OffloadAPI/CMakeLists.txt index 292ea1eb4852f..8bed11be4d85e 100644 --- a/offload/unittests/OffloadAPI/CMakeLists.txt +++ b/offload/unittests/OffloadAPI/CMakeLists.txt @@ -49,6 +49,7 @@ add_offload_unittest("platform" add_offload_unittest("program" program/olCreateProgram.cpp program/olIsValidBinary.cpp + program/olIterateCompatibleDevices.cpp program/olDestroyProgram.cpp) add_offload_unittest("queue" diff --git a/offload/unittests/OffloadAPI/program/olIterateCompatibleDevices.cpp b/offload/unittests/OffloadAPI/program/olIterateCompatibleDevices.cpp new file mode 100644 index 0000000000000..afa26c57d6728 --- /dev/null +++ b/offload/unittests/OffloadAPI/program/olIterateCompatibleDevices.cpp @@ -0,0 +1,73 @@ +//===------- Offload API tests - olIterateCompatibleDevices -------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include "../common/Fixtures.hpp" +#include <OffloadAPI.h> +#include <gtest/gtest.h> + +using olIterateCompatibleDevicesTest = OffloadDeviceTest; +OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(olIterateCompatibleDevicesTest); + +TEST_P(olIterateCompatibleDevicesTest, Success) { + std::unique_ptr<llvm::MemoryBuffer> DeviceBin; + ASSERT_TRUE(TestEnvironment::loadDeviceBinary("foo", Device, DeviceBin)); + ASSERT_GE(DeviceBin->getBufferSize(), 0lu); + + struct CallbackDataTy { + ol_device_handle_t ExpectedDevice; + bool Found = false; + } CallbackData{Device}; + + ASSERT_SUCCESS(olIterateCompatibleDevices( + DeviceBin->getBufferStart(), DeviceBin->getBufferSize(), + [](ol_device_handle_t D, void *UserData) { + auto *Data = static_cast<CallbackDataTy *>(UserData); + if (D == Data->ExpectedDevice) + Data->Found = true; + return true; + }, + &CallbackData)); + + ASSERT_TRUE(CallbackData.Found); +} + +TEST_P(olIterateCompatibleDevicesTest, SuccessStopIteration) { + std::unique_ptr<llvm::MemoryBuffer> DeviceBin; + ASSERT_TRUE(TestEnvironment::loadDeviceBinary("foo", Device, DeviceBin)); + ASSERT_GE(DeviceBin->getBufferSize(), 0lu); + + uint32_t CallCount = 0; + ASSERT_SUCCESS(olIterateCompatibleDevices( + DeviceBin->getBufferStart(), DeviceBin->getBufferSize(), + [](ol_device_handle_t, void *UserData) { + auto *Count = static_cast<uint32_t *>(UserData); + *Count += 1; + return false; + }, + &CallCount)); + + ASSERT_EQ(CallCount, 1u); +} + +TEST_P(olIterateCompatibleDevicesTest, EmptyBinary) { + std::unique_ptr<llvm::MemoryBuffer> DeviceBin; + ASSERT_TRUE(TestEnvironment::loadDeviceBinary("foo", Device, DeviceBin)); + ASSERT_GE(DeviceBin->getBufferSize(), 0lu); + + uint32_t CallCount = 0; + ASSERT_SUCCESS(olIterateCompatibleDevices( + DeviceBin->getBufferStart(), 0, + [](ol_device_handle_t, void *UserData) { + auto *Count = static_cast<uint32_t *>(UserData); + *Count += 1; + return true; + }, + &CallCount)); + + ASSERT_EQ(CallCount, 0u); +} _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
