| 28 | |
| 29 | namespace paddle::pybind { |
| 30 | void BindCustomDevicePy(py::module *m_ptr) { |
| 31 | auto &m = *m_ptr; |
| 32 | // Bind Methods |
| 33 | m.def("_get_device_min_chunk_size", [](const std::string &device_type) { |
| 34 | auto place = phi::CustomPlace(device_type); |
| 35 | return phi::DeviceManager::GetMinChunkSize(place); |
| 36 | }); |
| 37 | m.def( |
| 38 | "_get_device_total_memory", |
| 39 | [](const std::string &device_type, int device_id) { |
| 40 | auto place = phi::CustomPlace( |
| 41 | device_type, |
| 42 | device_id == -1 ? phi::DeviceManager::GetDevice(device_type) |
| 43 | : device_id); |
| 44 | size_t total = 0, free = 0; |
| 45 | phi::DeviceManager::MemoryStats(place, &total, &free); |
| 46 | return total; |
| 47 | }, |
| 48 | py::arg("device_type"), |
| 49 | py::arg("device_id") = -1); |
| 50 | m.def( |
| 51 | "_get_current_custom_device_stream", |
| 52 | [](const std::string &device_type, int device_id) { |
| 53 | #ifdef PADDLE_WITH_CUSTOM_DEVICE |
| 54 | auto place = phi::CustomPlace( |
| 55 | device_type, |
| 56 | device_id == -1 ? phi::DeviceManager::GetDevice(device_type) |
| 57 | : device_id); |
| 58 | |
| 59 | return static_cast<const phi::CustomContext *>( |
| 60 | phi::DeviceContextPool::Instance().Get(place)) |
| 61 | ->GetStream(); |
| 62 | #else |
| 63 | PADDLE_THROW(common::errors::Unavailable( |
| 64 | "Paddle is not compiled with CustomDevice. " |
| 65 | "Cannot visit _get_current_custom_device_stream.")); |
| 66 | #endif |
| 67 | }, |
| 68 | py::return_value_policy::reference, |
| 69 | py::arg("device_type"), |
| 70 | py::arg("device_id") = -1); |
| 71 | m.def( |
| 72 | "_set_current_custom_device_stream", |
| 73 | [](const std::string &device_type, |
| 74 | int device_id, |
| 75 | std::shared_ptr<phi::stream::Stream> stream) { |
| 76 | #ifdef PADDLE_WITH_CUSTOM_DEVICE |
| 77 | auto place = phi::CustomPlace( |
| 78 | device_type, |
| 79 | device_id == -1 ? phi::DeviceManager::GetDevice(device_type) |
| 80 | : device_id); |
| 81 | static_cast<phi::CustomContext *>( |
| 82 | phi::DeviceContextPool::Instance().Get(place)) |
| 83 | ->SetStream(stream); |
| 84 | return stream; |
| 85 | #else |
| 86 | PADDLE_THROW(common::errors::Unavailable( |
| 87 | "Paddle is not compiled with CustomDevice. " |
no test coverage detected