| 2773 | } |
| 2774 | |
| 2775 | void ExposeStream(py::module &m) { |
| 2776 | py::class_<dali::CUDAStreamLease>(m, "Stream") |
| 2777 | .def(py::init([](std::optional<int> device_id) { |
| 2778 | return std::make_unique<dali::CUDAStreamLease>( |
| 2779 | dali::CUDAStreamPool::instance().Get(device_id.value_or(-1))); |
| 2780 | })) |
| 2781 | // __cuda_stream__ protocol, as per cuda.core 0.3.0 |
| 2782 | .def("__cuda_stream__", [](dali::CUDAStreamLease &self) { |
| 2783 | // version, stream - version is 0 |
| 2784 | return std::make_tuple(0, reinterpret_cast<uintptr_t>(self.get())); |
| 2785 | }) |
| 2786 | .def_property_readonly("device_id", [](dali::CUDAStreamLease &self) { |
| 2787 | return self.device_id(); |
| 2788 | }) |
| 2789 | .def_property_readonly("handle", [](dali::CUDAStreamLease &self) { |
| 2790 | return reinterpret_cast<uintptr_t>(self.get()); |
| 2791 | }) |
| 2792 | .def("__repr__", [](dali::CUDAStreamLease &self) { |
| 2793 | return make_string("Stream(", self.get(), ")"); |
| 2794 | }) |
| 2795 | .def("__str__", [](dali::CUDAStreamLease &self) { |
| 2796 | return make_string("Stream(", self.get(), ")"); |
| 2797 | }); |
| 2798 | } |
| 2799 | |
| 2800 | /** Python wrapper for the Workspace */ |
| 2801 | class PyWorkspace : public Workspace { |
no test coverage detected