From 9ec8701fa206f3a43e7d2143f36d06404c3c41f0 Mon Sep 17 00:00:00 2001
From: Olivier BICHLER <olivier.bichler@cea.fr>
Date: Fri, 27 Jun 2025 15:42:53 +0200
Subject: [PATCH] Added PaddedInPlace_CP

---
 aidge_export_cpp/operators/Conv.py   |  7 +++++--
 aidge_export_cpp/operators/ConvDw.py |  7 +++++--
 aidge_export_cpp/operators/Pad.py    | 22 ++++++++++++++++++++++
 aidge_export_cpp/operators/Pool.py   | 25 +++++++++++++++++++------
 examples/export_ResNet18/resnet18.py |  4 +++-
 5 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/aidge_export_cpp/operators/Conv.py b/aidge_export_cpp/operators/Conv.py
index a813c2f..88fbf3b 100644
--- a/aidge_export_cpp/operators/Conv.py
+++ b/aidge_export_cpp/operators/Conv.py
@@ -1,6 +1,7 @@
 import aidge_core
 from aidge_core.export_utils import ExportNodeCpp, get_node_from_metaop
 from aidge_export_cpp import ROOT, ExportLibCpp, set_scaling_attributes
+from .Pad import PaddedInPlace_CP
 
 @ExportLibCpp.register("Conv1D",
     aidge_core.ImplSpec(
@@ -136,7 +137,8 @@ class QConv(Conv):
         [ # Output specifications
             aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)
         ],
-    ))
+    ),
+    PaddedInPlace_CP.default_model)
 class PadConv(QConv):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
@@ -178,7 +180,8 @@ class ConvAct(QConv):
         [ # Output specifications
             aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)
         ],
-    ))
+    ),
+    PaddedInPlace_CP.default_model)
 class PadConvAct(PadConv, ConvAct):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
diff --git a/aidge_export_cpp/operators/ConvDw.py b/aidge_export_cpp/operators/ConvDw.py
index a38030e..d0587f5 100644
--- a/aidge_export_cpp/operators/ConvDw.py
+++ b/aidge_export_cpp/operators/ConvDw.py
@@ -1,6 +1,7 @@
 import aidge_core
 from aidge_core.export_utils import ExportNodeCpp, get_node_from_metaop
 from aidge_export_cpp import ROOT, ExportLibCpp, set_scaling_attributes
+from .Pad import PaddedInPlace_CP
 
 @ExportLibCpp.register("ConvDepthWise2D",
     aidge_core.ImplSpec(
@@ -87,7 +88,8 @@ class QConvDw(ConvDw):
         [ # Output specifications
             aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)
         ],
-    ))
+    ),
+    PaddedInPlace_CP.default_model)
 class PadConvDw(QConvDw):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
@@ -129,7 +131,8 @@ class ConvDwAct(QConvDw):
         [ # Output specifications
             aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)
         ],
-    ))
+    ),
+    PaddedInPlace_CP.default_model)
 class PadConvDwAct(PadConvDw, ConvDwAct):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
diff --git a/aidge_export_cpp/operators/Pad.py b/aidge_export_cpp/operators/Pad.py
index f84f2cf..57fd86d 100644
--- a/aidge_export_cpp/operators/Pad.py
+++ b/aidge_export_cpp/operators/Pad.py
@@ -3,6 +3,28 @@ from aidge_core.export_utils import ExportNodeCpp
 from aidge_export_cpp import ROOT
 from aidge_export_cpp import ExportLibCpp
 
+# Consumer-Producer model to allow memory wrapping for Padded in-place operators
+# like PaddedConv ou PaddedMax/AvgPool
+class PaddedInPlace_CP(aidge_core.ProdConso):
+    def __init__(self, op: aidge_core.Operator):
+        aidge_core.ProdConso.__init__(self, op, False)
+
+    def default_model(op: aidge_core.Operator):
+        return PaddedInPlace_CP(op)
+
+    def get_nb_required_protected(self, input_idx):
+        if input_idx != 0:
+            return super().get_nb_required_protected(input_idx)
+
+        pad_node = self.get_operator().get_micro_graph().get_ordered_inputs()[0][0]
+        input = pad_node.get_operator().get_input(0)
+        if input:
+            output = pad_node.get_operator().get_output(0)
+            return aidge_core.Elts_t.data_elts(output.size() - input.size())
+        else:
+            return aidge_core.Elts_t.none_elts()
+
+
 @ExportLibCpp.register("Pad2D", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any)))
 class CppPad(ExportNodeCpp):
     def __init__(self, node, mem_info):
diff --git a/aidge_export_cpp/operators/Pool.py b/aidge_export_cpp/operators/Pool.py
index d6e9bfc..9c8df13 100644
--- a/aidge_export_cpp/operators/Pool.py
+++ b/aidge_export_cpp/operators/Pool.py
@@ -2,6 +2,7 @@ import aidge_core
 from aidge_core.export_utils import ExportNodeCpp, get_node_from_metaop
 from aidge_export_cpp import ROOT
 from aidge_export_cpp import ExportLibCpp
+from .Pad import PaddedInPlace_CP
 
 class Pool(ExportNodeCpp):
     def __init__(self, node, mem_info):
@@ -66,7 +67,9 @@ class MaxPool(Pool):
         self.attributes["stride_dims"] = PoolNode[0].get_operator().attr.stride_dims
 
 
-@ExportLibCpp.register_metaop(["PaddedMaxPooling2D", "PadMaxPool"], aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)))
+@ExportLibCpp.register_metaop(["PaddedMaxPooling2D", "PadMaxPool"],
+    aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)),
+    PaddedInPlace_CP.default_model)
 class PadMaxPool(MaxPool, PadPool):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
@@ -80,7 +83,9 @@ class MaxPoolAct(MaxPool, PoolAct):
         super().__init__(node, mem_info)
 
 
-@ExportLibCpp.register_metaop("PadMaxPoolAct", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)))
+@ExportLibCpp.register_metaop("PadMaxPoolAct",
+    aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)),
+    PaddedInPlace_CP.default_model)
 class PadMaxPoolAct(PadMaxPool, MaxPoolAct):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
@@ -100,7 +105,9 @@ class AvgPool(Pool):
         self.attributes["stride_dims"] = PoolNode[0].get_operator().attr.stride_dims
 
 
-@ExportLibCpp.register_metaop(["PaddedAvgPooling2D", "PadAvgPool"], aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)))
+@ExportLibCpp.register_metaop(["PaddedAvgPooling2D", "PadAvgPool"],
+    aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)),
+    PaddedInPlace_CP.default_model)
 class PadAvgPool(AvgPool, PadPool):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
@@ -114,7 +121,9 @@ class AvgPoolAct(AvgPool, PoolAct):
         super().__init__(node, mem_info)
 
 
-@ExportLibCpp.register_metaop("PadAvgPoolAct", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)))
+@ExportLibCpp.register_metaop("PadAvgPoolAct",
+    aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)),
+    PaddedInPlace_CP.default_model)
 class PadAvgPoolAct(PadAvgPool, AvgPoolAct):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
@@ -131,7 +140,9 @@ class GlobalAvgPool(Pool):
         self.attributes["kernel_dims"] = [self.attributes["in_width"][0], self.attributes["in_height"][0]]
 
 
-@ExportLibCpp.register_metaop("PadGlobalAvgPool", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)))
+@ExportLibCpp.register_metaop("PadGlobalAvgPool",
+    aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)),
+    PaddedInPlace_CP.default_model)
 class PadGlobalAvgPool(GlobalAvgPool, PadPool):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
@@ -145,7 +156,9 @@ class GlobalAvgPoolAct(GlobalAvgPool, PoolAct):
         super().__init__(node, mem_info)
 
 
-@ExportLibCpp.register_metaop("PadGlobalAvgPoolAct", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)))
+@ExportLibCpp.register_metaop("PadGlobalAvgPoolAct",
+    aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)),
+    PaddedInPlace_CP.default_model)
 class PadGlobalAvgPoolAct(PadGlobalAvgPool, GlobalAvgPoolAct):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
\ No newline at end of file
diff --git a/examples/export_ResNet18/resnet18.py b/examples/export_ResNet18/resnet18.py
index b27cafd..4cca5c8 100644
--- a/examples/export_ResNet18/resnet18.py
+++ b/examples/export_ResNet18/resnet18.py
@@ -42,6 +42,7 @@ supported_types = ["float32", "int8"]
 
 parser = argparse.ArgumentParser(description="Export the ResNet18 model with the aidge_export_cpp module.")
 parser.add_argument("--dev", action="store_true", help="Export in dev mode")
+parser.add_argument("--mem_wrap", action="store_true", help="Use memory wrapping")
 parser.add_argument("--no_cuda", action="store_true", help="Disable USE_CUDA usage to perform inferences and training.")
 parser.add_argument("--dtype", type=str, choices=supported_types, default="float32", help="Specify the targeted datatype : [int8, float32]")
 parser.add_argument("--aidge_cmp", action="store_true", help="Use aidge tensor results as reference.")
@@ -559,7 +560,8 @@ aidge_export_cpp.export(EXPORT_FOLDER,
                         labels = aidge_core.Tensor(labels[0]),
                         #inputs_tensor=inputs_tensor,
                         dev_mode = DEV_MODE,
-                        aidge_cmp = AIDGE_CMP)
+                        aidge_cmp = AIDGE_CMP,
+                        memory_manager_args = {"wrapping": True} if args.mem_wrap else {})
 
 print("\n### Compiling the export ###")
 try:
-- 
GitLab