Add more gpu-related python bindings (#5165)

Tencent · Dec 21, 2023 · a750600 · a750600
1 parent d30af29
commit a750600
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/python/src/main.cpp b/python/src/main.cpp
@@ -1277,15 +1277,21 @@ 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_<VulkanDevice>(m, "VulkanDevice")
     .def(py::init<int>(), py::arg("device_index") = 0)
     .def(
     "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(