c7e7ce7569
--------- Co-authored-by: Dmitry Razdoburdin <> Co-authored-by: Jiaming Yuan <jm.yuan@outlook.com>
52 lines
1.2 KiB
C++
52 lines
1.2 KiB
C++
/*!
|
|
* Copyright 2017-2023 by Contributors
|
|
* \file device_manager.h
|
|
*/
|
|
#ifndef PLUGIN_SYCL_DEVICE_MANAGER_H_
|
|
#define PLUGIN_SYCL_DEVICE_MANAGER_H_
|
|
|
|
#include <vector>
|
|
#include <mutex>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
|
|
#include <CL/sycl.hpp>
|
|
|
|
#pragma GCC diagnostic push
|
|
#pragma GCC diagnostic ignored "-Wtautological-constant-compare"
|
|
#pragma GCC diagnostic ignored "-W#pragma-messages"
|
|
#include "xgboost/context.h"
|
|
#pragma GCC diagnostic pop
|
|
|
|
namespace xgboost {
|
|
namespace sycl {
|
|
|
|
class DeviceManager {
|
|
public:
|
|
::sycl::queue GetQueue(const DeviceOrd& device_spec) const;
|
|
|
|
::sycl::device GetDevice(const DeviceOrd& device_spec) const;
|
|
|
|
private:
|
|
using QueueRegister_t = std::unordered_map<std::string, ::sycl::queue>;
|
|
constexpr static int kDefaultOrdinal = -1;
|
|
|
|
struct DeviceRegister {
|
|
std::vector<::sycl::device> devices;
|
|
std::vector<::sycl::device> cpu_devices;
|
|
std::vector<::sycl::device> gpu_devices;
|
|
};
|
|
|
|
QueueRegister_t& GetQueueRegister() const;
|
|
|
|
DeviceRegister& GetDevicesRegister() const;
|
|
|
|
mutable std::mutex queue_registering_mutex;
|
|
mutable std::mutex device_registering_mutex;
|
|
};
|
|
|
|
} // namespace sycl
|
|
} // namespace xgboost
|
|
|
|
#endif // PLUGIN_SYCL_DEVICE_MANAGER_H_
|