Ensure AOT passes all intermediary storages to function calls

[Mousius]

This iterates over the return storage IDs rather than just using the
first one to ensure all of them get passed to subsequent calls.

Fixes #9036

[PhilippvK]

The actual fix good to me! Thank you.

One aspect I would like to mention before this is getting merged:

This currently replaces the original test_byoc_microtvm test case by the one I come up with to reproduce the issue. I would suggest to add a second test case instead which uses multiple BYOC subgraphs. For this it would be helpful to replace the existing CcompilerAnnotator class with a generic one which is able to annotate any given model. Ideally the MergeCompilerRegions transformation should be enabled/diabled via pytest like this: @pytest.mark.parametrize("merge_compiler_regions", [False, True])

[Mousius]

@PhilippvK, argh! The test said it was a simple BYOC case, I didn't realise it was exercising multiple sub graphs! I've re-added it for now, I think there's probably better test cases but fixing the bug takes priority - are you ok with just re-adding it?

[PhilippvK]

@Mousius oops I just realized, that I am in fact only generating a single subgraph, but a multi-output one. I will then just follow up with another test case and an improved CcompilerAnnotator after this is merged.

[areusch]

@PhilippvK @Mousius so i'm not sure i'm seeing where we've added a multiple-return-value test case here. could you clarify for me? sorry to be dense :).

[PhilippvK]

@areusch The test case I came up with to reproduce to bug is quite messy because it does use manually written compiler_begin/compiler_end annotations because the existing CcompilerAnnotator can only create one specific graph. I would like to first make this class more generic to annotate any given model and then reimplement my test case for BYOC and multiple return values. But as this is out of scope of this "bugfix" I should follow up with the actual test case after this is merged.

[areusch]

@PhilippvK i would like to make sure we are exercising the code change here in test, though. it's not clear to me we are...could you clarify?

[Mousius]

@areusch I ensured the test failed before I made it pass 😸

[areusch]

ok got it now, makes sense!