diff --git a/python/src/main.cpp b/python/src/main.cpp index acf4308aa..0bbae2d78 100644 --- a/python/src/main.cpp +++ b/python/src/main.cpp @@ -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 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"); diff --git a/python/src/utils.hpp b/python/src/utils.hpp new file mode 100644 index 000000000..367871344 --- /dev/null +++ b/python/src/utils.hpp @@ -0,0 +1,26 @@ + +#include +#include + +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; +} +} +} diff --git a/python/test/test_kompute.py b/python/test/test_kompute.py index e1bcee940..d99dd4485 100644 --- a/python/test/test_kompute.py +++ b/python/test/test_kompute.py @@ -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] +