thek0tyara/llama-cpp-turboquant
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 4

Added list_devices function and separated new utils function

Browse source
This commit is contained in:
Alejandro Saucedo 2021-03-20 08:49:51 +00:00
parent d23ea95e64
commit a6457ee1ec
3 changed files with 49 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

25
python/src/main.cpp
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");