diff --git a/python_binding/pybind_core.cpp b/python_binding/pybind_core.cpp index f12ab25bf60fb32fb3b91a59997007fd2e266e5d..7b38c2d72d5f4b2eed8d8bbf9f41f47144b51060 100644 --- a/python_binding/pybind_core.cpp +++ b/python_binding/pybind_core.cpp @@ -73,6 +73,7 @@ void init_Recipes(py::module&); void init_GraphViewHelper(py::module&); void init_Scheduler(py::module&); +void init_MemoryManager(py::module&); void init_TensorUtils(py::module&); void init_Filler(py::module&); @@ -136,6 +137,7 @@ void init_Aidge(py::module& m) { init_Recipes(m); init_GraphViewHelper(m); init_Scheduler(m); + init_MemoryManager(m); init_TensorUtils(m); init_Filler(m); } diff --git a/python_binding/scheduler/pybind_MemoryManager.cpp b/python_binding/scheduler/pybind_MemoryManager.cpp new file mode 100644 index 0000000000000000000000000000000000000000..c560ff40cdfc5d4d50923a8f82b765058ed43973 --- /dev/null +++ b/python_binding/scheduler/pybind_MemoryManager.cpp @@ -0,0 +1,108 @@ +/******************************************************************************** + * Copyright (c) 2024 CEA-List + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0. + * + * SPDX-License-Identifier: EPL-2.0 + * + ********************************************************************************/ + +#include <pybind11/pybind11.h> +#include <pybind11/stl.h> + +#include "aidge/scheduler/MemoryManager.hpp" + +namespace py = pybind11; + +namespace Aidge { + +void init_MemoryManager(py::module& m) +{ + py::enum_<MemoryManager::OptimizeStrategy>(m, "OptimizeStrategy") + .value("None", MemoryManager::OptimizeStrategy::None) + .value("OptimizeMaxLifetimeMinSizeFirst", MemoryManager::OptimizeStrategy::OptimizeMaxLifetimeMinSizeFirst) + .value("OptimizeMaxLifetimeMaxSizeFirst", MemoryManager::OptimizeStrategy::OptimizeMaxLifetimeMaxSizeFirst) + .value("OptimizeMaxHoleMaxLifetimeFirst", MemoryManager::OptimizeStrategy::OptimizeMaxHoleMaxLifetimeFirst) + .export_values(); + + py::class_<MemoryManager::MemorySpace, std::shared_ptr<MemoryManager::MemorySpace>>(m, "MemorySpace") + .def(py::init<MemoryManager::Clock_T, unsigned int, unsigned int, std::set<std::shared_ptr<Node>> >(), py::arg("clock"), py::arg("offset"), py::arg("size"), py::arg("dependencies") = std::set<std::shared_ptr<Node>>()) + .def_readwrite("offset", &MemoryManager::MemorySpace::offset) + .def_readwrite("size", &MemoryManager::MemorySpace::size) + .def_readwrite("dependencies", &MemoryManager::MemorySpace::dependencies) + .def_readwrite("allocated", &MemoryManager::MemorySpace::allocated) + .def_readwrite("released", &MemoryManager::MemorySpace::released); + + py::class_<MemoryManager::MemoryPlane, std::shared_ptr<MemoryManager::MemoryPlane>>(m, "MemoryPlane") + .def(py::init<std::shared_ptr<MemoryManager::MemorySpace>, + MemoryManager::Clock_T, unsigned int, unsigned int, + unsigned int, unsigned int, unsigned int>(), + py::arg("memSpace"), py::arg("clock"), py::arg("offset"), + py::arg("size"), py::arg("stride"), py::arg("length"), py::arg("count")) + .def_readwrite("memSpace", &MemoryManager::MemoryPlane::memSpace) + .def_readwrite("allocated", &MemoryManager::MemoryPlane::allocated) + .def_readwrite("offset", &MemoryManager::MemoryPlane::offset) + .def_readwrite("size", &MemoryManager::MemoryPlane::size) + .def_readwrite("stride", &MemoryManager::MemoryPlane::stride) + .def_readwrite("length", &MemoryManager::MemoryPlane::length) + .def_readwrite("count", &MemoryManager::MemoryPlane::count) + .def("get_size", &MemoryManager::MemoryPlane::getSize) + .def("get_useful_size", &MemoryManager::MemoryPlane::getUsefulSize) + .def("get_contiguous_offset", &MemoryManager::MemoryPlane::getContiguousOffset) + .def("get_contiguous_size", &MemoryManager::MemoryPlane::getContiguousSize) + .def("get_wrapped_offset", &MemoryManager::MemoryPlane::getWrappedOffset) + .def("get_wrapped_size", &MemoryManager::MemoryPlane::getWrappedSize) + .def("get_final_offset", &MemoryManager::MemoryPlane::getFinalOffset) + .def("get_upper_offset", &MemoryManager::MemoryPlane::getUpperOffset) + .def("get_limit", &MemoryManager::MemoryPlane::getLimit); + + py::class_<MemoryManager::MaxLifetimeMinSizeFirst>(m, "MaxLifetimeMinSizeFirst") + .def(py::init<unsigned int>(), py::arg("max_lifetime")) + .def_readonly("max_lifetime", &MemoryManager::MaxLifetimeMinSizeFirst::maxLifetime) + .def("__call__", &MemoryManager::MaxLifetimeMinSizeFirst::operator(), py::arg("p0"), py::arg("p1")); + + py::class_<MemoryManager::MaxLifetimeMaxSizeFirst>(m, "MaxLifetimeMaxSizeFirst") + .def(py::init<unsigned int>(), py::arg("max_lifetime")) + .def_readonly("max_lifetime", &MemoryManager::MaxLifetimeMaxSizeFirst::maxLifetime) + .def("__call__", &MemoryManager::MaxLifetimeMaxSizeFirst::operator(), py::arg("p0"), py::arg("p1")); + + py::class_<MemoryManager::MaxHoleMaxLifetimeFirst>(m, "MaxHoleMaxLifetimeFirst") + .def(py::init<unsigned int, MemoryManager*>(), py::arg("max_lifetime"), py::arg("inst")) + .def_readonly("max_lifetime", &MemoryManager::MaxHoleMaxLifetimeFirst::maxLifetime) + .def_readwrite("inst", &MemoryManager::MaxHoleMaxLifetimeFirst::inst) + .def("__call__", &MemoryManager::MaxHoleMaxLifetimeFirst::operator(), py::arg("p0"), py::arg("p1")); + + py::class_<MemoryManager, std::shared_ptr<MemoryManager>>(m, "MemoryManager") + .def(py::init<>()) + .def("reserve", (std::shared_ptr<MemoryManager::MemorySpace> (MemoryManager::*)(unsigned int, const std::set<std::shared_ptr<Node>>&)) &MemoryManager::reserve, py::arg("size"), py::arg("dependencies") = std::set<std::shared_ptr<Node>>()) + .def("expand", &MemoryManager::expand, py::arg("memSpace"), py::arg("required_size")) + .def("allocate", (MemoryManager::MemoryPlane (MemoryManager::*)(unsigned int, const std::set<std::shared_ptr<Node>>&, unsigned int, unsigned int, unsigned int)) &MemoryManager::allocate, py::arg("size"), py::arg("dependencies") = std::set<std::shared_ptr<Node>>(), py::arg("stride") = 0, py::arg("length") = 1, py::arg("count") = 1) + .def("allocate", (unsigned int (MemoryManager::*)(const std::shared_ptr<Node>&, unsigned int, const std::set<std::shared_ptr<Node>>&, unsigned int, unsigned int, unsigned int)) &MemoryManager::allocate, py::arg("node"), py::arg("size"), py::arg("dependencies") = std::set<std::shared_ptr<Node>>(), py::arg("stride") = 0, py::arg("length") = 1, py::arg("count") = 1) + .def("is_wrap_around", &MemoryManager::isWrapAround, py::arg("memSpace"), py::arg("offset"), py::arg("size"), py::arg("stride") = 0, py::arg("length") = 1, py::arg("count") = 1) + .def("reallocate", (MemoryManager::MemoryPlane (MemoryManager::*)(std::shared_ptr<MemoryManager::MemorySpace>, unsigned int, unsigned int, bool, unsigned int, const std::set<std::shared_ptr<Node>>&, unsigned int, unsigned int, unsigned int)) &MemoryManager::reallocate, py::arg("memSpace"), py::arg("offset"), py::arg("size"), py::arg("wrap_around"), py::arg("extra_size") = 0, py::arg("additional_dependencies") = std::set<std::shared_ptr<Node>>(), py::arg("stride") = 0, py::arg("length") = 1, py::arg("count") = 1) + .def("reallocate", (MemoryManager::MemoryPlane (MemoryManager::*)(const MemoryManager::MemoryPlane&, unsigned int, unsigned int, bool, unsigned int, const std::set<std::shared_ptr<Node>>&, unsigned int, unsigned int, unsigned int)) &MemoryManager::reallocate, py::arg("memPlane"), py::arg("extra_offset"), py::arg("size"), py::arg("wrap_around"), py::arg("extra_size") = 0, py::arg("additional_dependencies") = std::set<std::shared_ptr<Node>>(), py::arg("stride") = 0, py::arg("length") = 1, py::arg("count") = 1) + .def("reallocate", (unsigned int (MemoryManager::*)(std::shared_ptr<MemoryManager::MemorySpace>, const std::shared_ptr<Node>&, unsigned int, unsigned int, bool, unsigned int, const std::set<std::shared_ptr<Node>>&, unsigned int, unsigned int, unsigned int)) &MemoryManager::reallocate, py::arg("memSpace"), py::arg("node"), py::arg("offset"), py::arg("size"), py::arg("wrap_around"), py::arg("extra_size") = 0, py::arg("additional_dependencies") = std::set<std::shared_ptr<Node>>(), py::arg("stride") = 0, py::arg("length") = 1, py::arg("count") = 1) + .def("reallocate", (unsigned int (MemoryManager::*)(const MemoryManager::MemoryPlane&, const std::shared_ptr<Node>&, unsigned int, unsigned int, bool, unsigned int, const std::set<std::shared_ptr<Node>>&, unsigned int, unsigned int, unsigned int)) &MemoryManager::reallocate, py::arg("memPlane"), py::arg("node"), py::arg("extra_offset"), py::arg("size"), py::arg("wrap_around"), py::arg("extra_size") = 0, py::arg("additional_dependencies") = std::set<std::shared_ptr<Node>>(), py::arg("stride") = 0, py::arg("length") = 1, py::arg("count") = 1) + .def("release", (unsigned int (MemoryManager::*)(std::shared_ptr<MemoryManager::MemorySpace>)) &MemoryManager::release, py::arg("memSpace")) + .def("release", (unsigned int (MemoryManager::*)(const std::shared_ptr<Node>&)) &MemoryManager::release, py::arg("node")) + .def("releaseDependencies", &MemoryManager::releaseDependencies, py::arg("node")) + .def("optimize", &MemoryManager::optimize, py::arg("strategy")) + .def("get_offset", &MemoryManager::getOffset, py::arg("node"), py::arg("plane") = 0) + .def("get_size", (unsigned int (MemoryManager::*)(const std::shared_ptr<Node>&, unsigned int) const) &MemoryManager::getSize, py::arg("node"), py::arg("plane")) + .def("get_size", (unsigned int (MemoryManager::*)(const std::shared_ptr<Node>&) const) &MemoryManager::getSize, py::arg("node")) + .def("get_peak_usage", &MemoryManager::getPeakUsage) + .def("get_max_lifetime", &MemoryManager::getMaxLifetime) + .def("get_planes", (const std::vector<MemoryManager::MemoryPlane>& (MemoryManager::*)(const std::shared_ptr<Node>&) const) &MemoryManager::getPlanes, py::arg("node")) + .def("get_planes", (const MemoryManager::MemMap_T& (MemoryManager::*)() const) &MemoryManager::getPlanes) + .def("get_planes", (MemoryManager::MemMap_T (MemoryManager::*)(std::shared_ptr<MemoryManager::MemorySpace>) const) &MemoryManager::getPlanes, py::arg("memSpace")) + .def("get_nb_planes", (unsigned int (MemoryManager::*)(const std::shared_ptr<Node>&) const) &MemoryManager::getNbPlanes, py::arg("node")) + .def("get_nb_planes", (unsigned int (MemoryManager::*)(std::shared_ptr<MemoryManager::MemorySpace>) const) &MemoryManager::getNbPlanes, py::arg("memSpace")) + .def("get_current_tick", &MemoryManager::getCurrentTick) + .def("tick", &MemoryManager::tick) + .def("log", &MemoryManager::log, py::arg("file_name")) + ; +} + +} // Aidge diff --git a/python_binding/scheduler/pybind_Scheduler.cpp b/python_binding/scheduler/pybind_Scheduler.cpp index c0966e54d4f025a607aa9763a3657de5b39d2ff4..d9b626bab8ca78c2f6662f5667a423b2659dd6ed 100644 --- a/python_binding/scheduler/pybind_Scheduler.cpp +++ b/python_binding/scheduler/pybind_Scheduler.cpp @@ -11,6 +11,7 @@ #include <pybind11/pybind11.h> #include <pybind11/stl.h> +#include "aidge/scheduler/MemoryManager.hpp" #include "aidge/scheduler/Scheduler.hpp" #include "aidge/scheduler/SequentialScheduler.hpp" #include "aidge/scheduler/ParallelScheduler.hpp" @@ -26,6 +27,7 @@ void init_Scheduler(py::module& m){ .def("resetScheduling", &Scheduler::resetScheduling) .def("generate_scheduling", &Scheduler::generateScheduling) .def("get_static_scheduling", &Scheduler::getStaticScheduling, py::arg("step") = 0) + .def("generate_memory", &Scheduler::generateMemory, py::arg("inc_producers") = false, py::arg("wrap_around_buffer") = false) ; py::class_<SequentialScheduler, std::shared_ptr<SequentialScheduler>, Scheduler>(m, "SequentialScheduler")