Merge pull request #195 from EthicalML/185_list_devices

Add function to list physical devices
This commit is contained in:
Alejandro Saucedo 2021-03-20 18:59:34 +00:00 committed by GitHub
commit b0df305daf
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 89 additions and 19 deletions

View file

@ -6,10 +6,10 @@
#include "fmt/ranges.h"
#include "utils.hpp"
#include "docstrings.hpp"
namespace py = pybind11;
using namespace pybind11::literals; // for the `_a` literal
//used in Core.hpp
py::object kp_debug, kp_info, kp_warning, kp_error;
@ -220,21 +220,18 @@ PYBIND11_MODULE(kp, m) {
py::arg("workgroup") = kp::Workgroup(),
py::arg("spec_consts") = kp::Constants(),
py::arg("push_consts") = kp::Constants())
.def("list_devices", [](kp::Manager& self){
const std::vector<vk::PhysicalDevice> devices = self.listDevices();
py::list list;
for (const vk::PhysicalDevice& device : devices) {
list.append(kp::py::vkPropertiesToDict(device.getProperties()));
}
return list;
}, "Return a dict containing information about the device")
.def("get_device_properties", [](kp::Manager& self){
const auto properties = self.getDeviceProperties();
py::dict py_props(
"device_name"_a = std::string(properties.deviceName.data()),
"max_work_group_count"_a = py::make_tuple(properties.limits.maxComputeWorkGroupCount[0],
properties.limits.maxComputeWorkGroupCount[1],
properties.limits.maxComputeWorkGroupCount[2]),
"max_work_group_invocations"_a = properties.limits.maxComputeWorkGroupInvocations,
"max_work_group_size"_a = py::make_tuple(properties.limits.maxComputeWorkGroupSize[0],
properties.limits.maxComputeWorkGroupSize[1],
properties.limits.maxComputeWorkGroupSize[2]),
"timestamps_supported"_a = (bool)properties.limits.timestampComputeAndGraphics
);
const vk::PhysicalDeviceProperties properties = self.getDeviceProperties();
return py_props;
return kp::py::vkPropertiesToDict(properties);
}, "Return a dict containing information about the device");

26
python/src/utils.hpp Normal file
View file

@ -0,0 +1,26 @@
#include <kompute/Kompute.hpp>
#include <pybind11/pybind11.h>
using namespace pybind11::literals; // for the `_a` literal
namespace kp {
namespace py {
static pybind11::dict vkPropertiesToDict(const vk::PhysicalDeviceProperties& properties) {
pybind11::dict pyDict(
"device_name"_a = std::string(properties.deviceName.data()),
"max_work_group_count"_a = pybind11::make_tuple(properties.limits.maxComputeWorkGroupCount[0],
properties.limits.maxComputeWorkGroupCount[1],
properties.limits.maxComputeWorkGroupCount[2]),
"max_work_group_invocations"_a = properties.limits.maxComputeWorkGroupInvocations,
"max_work_group_size"_a = pybind11::make_tuple(properties.limits.maxComputeWorkGroupSize[0],
properties.limits.maxComputeWorkGroupSize[1],
properties.limits.maxComputeWorkGroupSize[2]),
"timestamps_supported"_a = (bool)properties.limits.timestampComputeAndGraphics
);
return pyDict;
}
}
}

View file

@ -227,3 +227,15 @@ def test_workgroup():
assert np.all(tensor_a.data() == np.stack([np.arange(16)]*8, axis=1).ravel())
assert np.all(tensor_b.data() == np.stack([np.arange(8)]*16, axis=0).ravel())
def test_mgr_utils():
mgr = kp.Manager()
props = mgr.get_device_properties()
assert "device_name" in props
devices = mgr.list_devices()
assert len(devices) == 1
assert "device_name" in devices[0]

View file

@ -2156,10 +2156,19 @@ class Manager
void clear();
/**
* Return a struct containing information about the device.
* Information about the current device.
*
* @return vk::PhysicalDeviceProperties containing information about the device
**/
vk::PhysicalDeviceProperties getDeviceProperties() const;
/**
* List the devices available in the current vulkan instance.
*
* @return vector of physical devices containing their respective properties
**/
std::vector<vk::PhysicalDevice> listDevices() const;
private:
// -------------- OPTIONALLY OWNED RESOURCES
std::shared_ptr<vk::Instance> mInstance = nullptr;

View file

@ -51,8 +51,10 @@ endif()
target_include_directories(
kompute PUBLIC
${CMAKE_CURRENT_SOURCE_DIR}/include
${PROJECT_SOURCE_DIR}/single_include
$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
$<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}/single_include>
$<INSTALL_INTERFACE:include>
$<INSTALL_INTERFACE:single_include>
)
if(NOT KOMPUTE_OPT_ANDOID_BUILD)

View file

@ -453,4 +453,10 @@ Manager::getDeviceProperties() const
return this->mPhysicalDevice->getProperties();
}
std::vector<vk::PhysicalDevice>
Manager::listDevices() const
{
return this->mInstance->enumeratePhysicalDevices();
}
}

View file

@ -154,10 +154,20 @@ class Manager
void clear();
/**
* Return a struct containing information about the device.
* Information about the current device.
*
* @return vk::PhysicalDeviceProperties containing information about the device
**/
vk::PhysicalDeviceProperties getDeviceProperties() const;
/**
* List the devices available in the current vulkan instance.
*
* @return vector of physical devices containing their respective properties
**/
std::vector<vk::PhysicalDevice> listDevices() const;
private:
// -------------- OPTIONALLY OWNED RESOURCES
std::shared_ptr<vk::Instance> mInstance = nullptr;

View file

@ -66,6 +66,14 @@ TEST(TestManager, TestMultipleSequences)
TEST(TestManager, TestDeviceProperties)
{
kp::Manager mgr;
const auto properties = mgr.getDeviceProperties();
const vk::PhysicalDeviceProperties properties = mgr.getDeviceProperties();
EXPECT_GT(properties.deviceName.size(), 0);
}
TEST(TestManager, TestListDevices)
{
kp::Manager mgr;
const std::vector<vk::PhysicalDevice> devices = mgr.listDevices();
EXPECT_GT(devices.size(), 0);
EXPECT_GT(devices[0].getProperties().deviceName.size(), 0);
}