Fixed in-place CP model for Conv and Pool

aidge_export_cpp/operators/Conv.py

 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
+
+# Consumer-Producer model to allow memory wrapping for Conv/PaddedConv
+# (and Pool/PaddedPool), keeping one input line margin in NHWC data format 
+# (one input line = W*C)
+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)
+
+        input = self.get_operator().get_input(0)
+        if not input:
+            return aidge_core.Elts_t.none_elts()
+
+        # Non-Padded case: margin = one input line
+        margin = 1
+        if not self.get_operator().is_atomic():
+            # Padded case: margin = (padding_y / stride_y) input lines
+            sub_graph = self.get_operator().get_micro_graph().clone()
+            aidge_core.expand_metaops(sub_graph, True)
+
+            padding_y = 0
+            stride_y = 1
+            for node in sub_graph.get_nodes():
+                if hasattr(node.get_operator().attr, 'stride_dims'):
+                    if len(node.get_operator().attr.stride_dims) > 1:
+                        stride_y = node.get_operator().attr.stride_dims[0]
+                elif hasattr(node.get_operator().attr, 'begin_end_borders'):
+                    if len(node.get_operator().attr.begin_end_borders) > 2:
+                        padding_y = node.get_operator().attr.begin_end_borders[0]
+
+            margin += padding_y // stride_y
+
+        if len(input.dims()) == 4:
+            # 2D: one input line = W*C
+            margin *= input.dims()[2] * input.dims()[3]
+        else:
+            # 1D: one input line = C
+            margin *= input.dims()[2]
+
+        return aidge_core.Elts_t.data_elts(margin)
+
 
 @ExportLibCpp.register("Conv1D",
     aidge_core.ImplSpec(
@@ -14,7 +60,7 @@ from .Pad import PaddedInPlace_CP
             aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nwc)
         ],
     ),
-    aidge_core.ProdConso.in_place_model)
+    PaddedInPlace_CP.default_model)
 class Conv1D(ExportNodeCpp):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
@@ -64,7 +110,7 @@ class Conv1D(ExportNodeCpp):
             aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)
         ],
     ),
-    aidge_core.ProdConso.in_place_model)
+    PaddedInPlace_CP.default_model)
 class Conv2D(ExportNodeCpp):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
@@ -114,7 +160,7 @@ class Conv2D(ExportNodeCpp):
             aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)
         ],
     ),
-    aidge_core.ProdConso.in_place_model)
+    PaddedInPlace_CP.default_model)
 class QConv(Conv2D):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
@@ -159,7 +205,7 @@ class PadConv(QConv):
             aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)
         ],
     ),
-    aidge_core.ProdConso.in_place_model)
+    PaddedInPlace_CP.default_model)
 class ConvAct(QConv):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)

aidge_export_cpp/operators/ConvDw.py

 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
+from .Conv import PaddedInPlace_CP
 
 @ExportLibCpp.register("ConvDepthWise2D",
     aidge_core.ImplSpec(
@@ -14,7 +14,7 @@ from .Pad import PaddedInPlace_CP
             aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)
         ],
     ),
-    aidge_core.ProdConso.in_place_model)
+    PaddedInPlace_CP.default_model)
 class ConvDw(ExportNodeCpp):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
@@ -65,7 +65,7 @@ class ConvDw(ExportNodeCpp):
             aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)
         ],
     ),
-    aidge_core.ProdConso.in_place_model)
+    PaddedInPlace_CP.default_model)
 class QConvDw(ConvDw):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
@@ -110,7 +110,7 @@ class PadConvDw(QConvDw):
             aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)
         ],
     ),
-    aidge_core.ProdConso.in_place_model)
+    PaddedInPlace_CP.default_model)
 class ConvDwAct(QConvDw):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)

aidge_export_cpp/operators/Pad.py

@@ -3,20 +3,19 @@ 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):
+# Consumer-Producer model to allow memory wrapping for Pad in-place operator
+class PadInPlace_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)
+        return PadInPlace_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]
+        pad_node = self.get_operator()
         input = pad_node.get_operator().get_input(0)
         if input:
             output = pad_node.get_operator().get_output(0)
@@ -25,7 +24,9 @@ class PaddedInPlace_CP(aidge_core.ProdConso):
             return aidge_core.Elts_t.none_elts()
 
 
-@ExportLibCpp.register("Pad2D", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any)))
+@ExportLibCpp.register("Pad2D",
+    aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any)),
+    PadInPlace_CP.default_model)
 class CppPad(ExportNodeCpp):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)

aidge_export_cpp/operators/Pool.py

@@ -2,7 +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
+from .Conv import PaddedInPlace_CP
 
 class Pool(ExportNodeCpp):
     def __init__(self, node, mem_info):
@@ -55,7 +55,7 @@ class PoolAct(Pool):
 
 @ExportLibCpp.register("MaxPooling2D",
     aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)),
-    aidge_core.ProdConso.in_place_model)
+    PaddedInPlace_CP.default_model)
 class MaxPool(Pool):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
@@ -77,7 +77,7 @@ class PadMaxPool(MaxPool, PadPool):
 
 @ExportLibCpp.register_metaop("MaxPoolAct",
     aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)),
-    aidge_core.ProdConso.in_place_model)
+    PaddedInPlace_CP.default_model)
 class MaxPoolAct(MaxPool, PoolAct):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
@@ -93,7 +93,7 @@ class PadMaxPoolAct(PadMaxPool, MaxPoolAct):
 
 @ExportLibCpp.register("AvgPooling2D",
     aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)),
-    aidge_core.ProdConso.in_place_model)
+    PaddedInPlace_CP.default_model)
 class AvgPool(Pool):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
@@ -115,7 +115,7 @@ class PadAvgPool(AvgPool, PadPool):
 
 @ExportLibCpp.register_metaop("AvgPoolAct",
     aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any, aidge_core.dformat.nhwc)),
-    aidge_core.ProdConso.in_place_model)
+    PaddedInPlace_CP.default_model)
 class AvgPoolAct(AvgPool, PoolAct):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)