How to change scheduler behavior for multi-layer spatial tiling ?

• aidge_core: 76641b90544401ec3c9901fc191d72b4ccffd9f4

Problem description :

I want to control the memory peak of the cnn-based encoder of my model. For that, I implemented the multi-layer spatial tiling from AIDGE documentation : https://eclipse.dev/aidge/source/UserGuide/optimizeGraph.html.

See my diagram bellow : memory peak should be reduced if the branches are scheduled in sequence (yellow). My problem is that the scheduler computes the tiled branches in parallel (red)! That is the reason of my issue : how to change the scheduler behaviour so it resolves my branches in a sequential and not parallel order ?

Bellow diagram shows in a simplifed example how the tiling should reduce the mem_peak and how the actual scheduling does not reduce it :

Problem description in detail :

here is my model : classifier_model.onnx , random numpy input : test_input.npy , code : tiling_issue.py

what it does is :

1. Import the model from onnx to aidge graph
2. Fold my conv1d+ReLU nodes into a metaOp "encoder"
3. Compute receptive field to slice the input properly
4. Create a new encoder_tiled graph : slice the input into 2 slices + connect each slice to an "encoder" (same weights) + concatenate each encoder output.
5. Create the scheduler and validate that graph and encoder_tiled provide same results
6. visualize the node ordering and mem_peak : slices are scheduled in parallel

Problem vizualisation :

Mermaid diagram for the initial model

Mermaid diagram for the tiled model

Mermaid diagram for tiled static scheduling

memory_info for initial model :

memory_info for tiled encoder :