From a7506008c830d77c5e3baa35a6b9935c75d2f93a Mon Sep 17 00:00:00 2001 From: Joey Ballentine <34788790+joeyballentine@users.noreply.github.com> Date: Thu, 21 Dec 2023 07:50:08 +0100 Subject: [PATCH] Add more gpu-related python bindings (#5165) --- python/src/main.cpp | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/python/src/main.cpp b/python/src/main.cpp index 97cf2015cf6..a7ed0528c6a 100644 --- a/python/src/main.cpp +++ b/python/src/main.cpp @@ -1277,7 +1277,8 @@ PYBIND11_MODULE(ncnn, m) .def("pipeline_cache_uuid", [](GpuInfo& gpuinfo) { return py::memoryview::from_buffer(gpuinfo.pipeline_cache_uuid(), {VK_UUID_SIZE}, {sizeof(uint8_t) * VK_UUID_SIZE}); }) - .def("type", &GpuInfo::type); + .def("type", &GpuInfo::type) + .def("device_name", &GpuInfo::device_name); py::class_(m, "VulkanDevice") .def(py::init(), py::arg("device_index") = 0) @@ -1285,7 +1286,12 @@ PYBIND11_MODULE(ncnn, m) "info", [](VulkanDevice& dev) { return &dev.info; }, - py::return_value_policy::reference_internal); + py::return_value_policy::reference_internal) + .def("acquire_blob_allocator", &VulkanDevice::acquire_blob_allocator) + .def("reclaim_blob_allocator", &VulkanDevice::reclaim_blob_allocator, py::arg("vkallocator")) + .def("acquire_staging_allocator", &VulkanDevice::acquire_staging_allocator) + .def("reclaim_staging_allocator", &VulkanDevice::reclaim_staging_allocator, py::arg("vkallocator")) + .def("get_heap_budget", &VulkanDevice::get_heap_budget); #endif // NCNN_VULKAN m.doc() = R"pbdoc(