| 2875 | } |
| 2876 | |
| 2877 | void ExposeWorkspace(py::module &m) { |
| 2878 | // Expose PyWorkspace as a private class within the backend_impl module |
| 2879 | py::class_<PyWorkspace>(m, "_Workspace") |
| 2880 | .def(py::init([](std::shared_ptr<ThreadPool> thread_pool, py::object stream) { |
| 2881 | auto ws = std::make_unique<PyWorkspace>(); |
| 2882 | ws->SetThreadPool(std::move(thread_pool)); |
| 2883 | if (!stream.is_none()) { |
| 2884 | ws->SetStream(stream); |
| 2885 | } |
| 2886 | return ws; |
| 2887 | }), "thread_pool"_a = nullptr, "stream"_a = py::none()) |
| 2888 | .def("SetStream", &PyWorkspace::SetStream, "cuda_stream"_a = py::none()) |
| 2889 | .def("AddInput", [](PyWorkspace &self, const TensorList<CPUBackend> &tl) { |
| 2890 | self.AddInput(CloneTL(tl)); |
| 2891 | }) |
| 2892 | .def("AddInput", [](PyWorkspace &self, const TensorList<GPUBackend> &tl) { |
| 2893 | self.AddInput(CloneTL(tl)); |
| 2894 | }) |
| 2895 | .def("AddArgumentInput", []( |
| 2896 | PyWorkspace &self, |
| 2897 | std::string name, |
| 2898 | const TensorList<CPUBackend> &tl) { |
| 2899 | self.AddArgumentInput(std::move(name), CloneTL(tl)); |
| 2900 | }) |
| 2901 | .def("AddOutput", [](PyWorkspace &self, const TensorList<CPUBackend> &tl) { |
| 2902 | self.AddOutput(CloneTL(tl)); |
| 2903 | }) |
| 2904 | .def("AddOutput", [](PyWorkspace &self, const TensorList<GPUBackend> &tl) { |
| 2905 | self.AddOutput(CloneTL(tl)); |
| 2906 | }) |
| 2907 | .def("AddEmptyOutput", [](PyWorkspace &self, std::string_view device) { |
| 2908 | auto storage_device = ParseStorageDevice(device); |
| 2909 | if (storage_device == StorageDevice::CPU) { |
| 2910 | self.AddOutput(std::make_shared<TensorList<CPUBackend>>()); |
| 2911 | } else { |
| 2912 | assert(storage_device == StorageDevice::GPU); |
| 2913 | self.AddOutput(std::make_shared<TensorList<GPUBackend>>()); |
| 2914 | } |
| 2915 | }, "device"_a) |
| 2916 | .def("SetThreadPool", [](PyWorkspace &self, std::shared_ptr<ThreadPool> thread_pool) { |
| 2917 | self.SetThreadPool(std::move(thread_pool)); |
| 2918 | }, "thread_pool"_a) |
| 2919 | .def("GetOutputs", [](PyWorkspace &self) { |
| 2920 | py::tuple ret(self.NumOutput()); |
| 2921 | for (int i = 0; i < self.NumOutput(); i++) { |
| 2922 | if (self.OutputIsType<CPUBackend>(i)) { |
| 2923 | ret[i] = self.OutputPtr<CPUBackend>(i); |
| 2924 | } else { |
| 2925 | ret[i] = self.OutputPtr<GPUBackend>(i); |
| 2926 | } |
| 2927 | } |
| 2928 | return ret; |
| 2929 | }, py::return_value_policy::take_ownership); |
| 2930 | } |
| 2931 | |
| 2932 | void SetupAndRun(OperatorBase &self, Workspace &ws, std::optional<int> batch_size) { |
| 2933 | DomainTimeRange setup_and_run_tr("SetupAndRun " + GetOpDisplayName(self.GetSpec(), true), |
no test coverage detected