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/recipes/pybind_Recipes.cpp b/python_binding/recipes/pybind_Recipes.cpp
index f122c411618ce28a641fd46ee568f99cc48e9f58..b85d1c41ed90a5774a9b24062dfda4186c2294d5 100644
--- a/python_binding/recipes/pybind_Recipes.cpp
+++ b/python_binding/recipes/pybind_Recipes.cpp
@@ -21,66 +21,70 @@
namespace py = pybind11;
namespace Aidge {
-void init_Recipes(py::module &m) {
+void init_Recipes(py::module &m)
+{
m.def("fuse_mul_add", static_cast<void(*)(std::shared_ptr<GraphView>)>(fuseMulAdd), py::arg("graph_view"), R"mydelimiter(
- Recipie to Fuse MatMul and Add operators into an :py:class:`aidge_core.FC` operator.
+ Recipe to Fuse MatMul and Add operators into an :py:class:`aidge_core.FC` operator.
- :param graph_view: Graph view on which we want to apply the recipie
+ :param graph_view: Graph view on which we want to apply the recipe
:type graph_view: :py:class:`aidge_core.GraphView`
)mydelimiter");
// m.def("fuse_mul_add", static_cast<void(*)(std::set<std::shared_ptr<Node>>)>(fuseMulAdd), py::arg("nodes"), R"mydelimiter(
- // Recipie to Fuse MatMul and Add operators into an :py:class:`aidge_core.FC` operator.
+ // recipe to Fuse MatMul and Add operators into an :py:class:`aidge_core.FC` operator.
// :param nodes: The MatMul and Add nodes to fuse.
// :type nodes: list of :py:class:`aidge_core.Node`
// )mydelimiter");
m.def("remove_dropout",static_cast<void(*)(std::shared_ptr<GraphView>)>(removeDropout), py::arg("graph_view"), R"mydelimiter(
- Recipie to remove a dropout operator.
+ Recipe to remove a dropout operator.
- :param graph_view: Graph view on which we want to apply the recipie
+ :param graph_view: Graph view on which we want to apply the recipe
:type graph_view: :py:class:`aidge_core.GraphView`
)mydelimiter");
m.def("remove_flatten", static_cast<void(*)(std::shared_ptr<GraphView>)>(removeFlatten), py::arg("graph_view"), R"mydelimiter(
- Recipie to remove a flatten operator.
+ Recipe to remove a flatten operator.
- :param graph_view: Graph view on which we want to apply the recipie
+ :param graph_view: Graph view on which we want to apply the recipe
:type graph_view: :py:class:`aidge_core.GraphView`
)mydelimiter");
// m.def("remove_flatten", static_cast<void(*)(std::set<std::shared_ptr<Node>>)>(removeFlatten), py::arg("nodes"), R"mydelimiter(
- // Recipie to remove a flatten operator.
+ // Recipe to remove a flatten operator.
// :param nodes: The flatten operator to remove.
// :type nodes: list of :py:class:`aidge_core.Node`
// )mydelimiter");
// m.def("fuse_mul_add", static_cast<void(*)(std::set<std::shared_ptr<Node>>)>(fuseMulAdd), py::arg("nodes"), R"mydelimiter(
- // Recipie to Fuse MatMul and Add operators into an :py:class:`aidge_core.FC` operator.
+ // Recipe to Fuse MatMul and Add operators into an :py:class:`aidge_core.FC` operator.
// :param nodes: The MatMul and Add nodes to fuse.
// :type nodes: list of :py:class:`aidge_core.Node`
// )mydelimiter");
m.def("fuse_batchnorm", static_cast<void(*)(std::shared_ptr<GraphView>)>(fuseBatchNorm), py::arg("graph_view"), R"mydelimiter(
- Recipie to remove a flatten operator.
+ Recipe to remove a flatten operator.
- :param graph_view: Graph view on which we want to apply the recipie
+ :param graph_view: Graph view on which we want to apply the recipe
:type graph_view: :py:class:`aidge_core.GraphView`
)mydelimiter");
- m.def("get_conv_horizontal_tiling", static_cast<std::set<std::shared_ptr<Node>>(*)(const std::shared_ptr<Node>&, const DimIdx_t, const std::size_t)>(getConvHorizontalTiling),
+ m.def("get_conv_horizontal_tiling", static_cast<std::set<std::shared_ptr<Node>>(*)(const std::shared_ptr<Node>&, const DimIdx_t, const std::size_t)>(getConvHorizontalTiling),
py::arg("node"), py::arg("axis"), py::arg("nb_slices"));
// m.def("fuse_batchnorm", static_cast<void(*)(std::set<std::shared_ptr<Node>>)>(fuseBatchNorm), py::arg("nodes"), R"mydelimiter(
- // Recipie to remove a flatten operator.
+ // recipe to remove a flatten operator.
// :param nodes: The flatten operator to remove.
// :type nodes: list of :py:class:`aidge_core.Node`
// )mydelimiter");
+
+ m.def("expand_metaops", static_cast<void(*)(std::shared_ptr<GraphView>, bool)>(expandMetaOps), py::arg("graph_view"), py::arg("recursive") = false);
}
+
} // namespace Aidge
diff --git a/python_binding/scheduler/pybind_MemoryManager.cpp b/python_binding/scheduler/pybind_MemoryManager.cpp
new file mode 100644
index 0000000000000000000000000000000000000000..0f18db405bec0aee9637f2e5f2ecc7b71e502cc5
--- /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("mem_space"), py::arg("clock"), py::arg("offset"),
+ py::arg("size"), py::arg("stride"), py::arg("length"), py::arg("count"))
+ .def_readwrite("mem_space", &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("mem_space"), 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("mem_space"), 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("mem_space"), 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("mem_space"), 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("mem_plane"), 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("mem_space"))
+ .def("release", (unsigned int (MemoryManager::*)(const std::shared_ptr<Node>&)) &MemoryManager::release, py::arg("node"))
+ .def("release_dependencies", &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("mem_space"))
+ .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("mem_space"))
+ .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..3f763c8ff0717fb07c1b6c1f85b6aba06c1dc8f1 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"
@@ -22,10 +23,12 @@ namespace Aidge {
void init_Scheduler(py::module& m){
py::class_<Scheduler, std::shared_ptr<Scheduler>>(m, "Scheduler")
.def(py::init<std::shared_ptr<GraphView>&>(), py::arg("graph_view"))
+ .def("graph_view", &Scheduler::graphView)
.def("save_scheduling_diagram", &Scheduler::saveSchedulingDiagram, py::arg("file_name"))
.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")
diff --git a/src/scheduler/Scheduler.cpp b/src/scheduler/Scheduler.cpp
index 4e3f9978837120bd01a3de2cfe2d22e33f9d7828..d85bd7cd48c0e537cdf5dc51524773e0de9519a2 100644
--- a/src/scheduler/Scheduler.cpp
+++ b/src/scheduler/Scheduler.cpp
@@ -491,17 +491,17 @@ Aidge::MemoryManager Aidge::Scheduler::generateMemory(bool incProducers, bool wr
const MemoryManager::MemoryPlane& memPlane
= (wrapAroundBuffer && wrapAroundSize > 0)
? (*wrapAroundMemPlane[outputIdx]) :
- memManager.allocate(requiredSize.data, childs, stride, length, count);
+ memManager.allocate(size, childs, stride, length, count);
if (wrapAroundBuffer && wrapAroundSize > 0) {
memManager.reallocate(memPlane,
node, 0,
- requiredSize.data, true, wrapAroundExtra, childs, stride, length, count);
+ size, true, wrapAroundExtra, childs, stride, length, count);
}
else {
memManager.reallocate(memPlane.memSpace,
node, memPlane.offset,
- requiredSize.data, false, 0, childs, stride, length, count);
+ size, false, 0, childs, stride, length, count);
}
}