MoveLinearPastEltWiseAdd() transformation bug

[penrosed]

Quick summary

When applying the MoveLinearPastEltWiseAdd() transformation on a graph where there are two identical linear producers of an element-wise add, if one of the linear nodes is a fork, the graph gets mangled.

Details

Steps to Reproduce

1. Clone the FINN repository
2. Checkout the dev branch, with commit hash: fd72d48
3. Download 'resnet18-4w4a_tidied.onnx'
4. Create a 'build.py' script in the same directory as 'resnet18-4w4a_tidied.onnx' which applies the MoveLinearPastEltWiseAdd() transformation on 'resnet18-4w4a.onnx', then saves the modified ONNX model. Here's an example 'build.py':

from sys import argv
from qonnx.core.modelwrapper import ModelWrapper
from finn.transformation.streamline.reorder import MoveLinearPastEltwiseAdd

# Set the filename that we look for our ONNX model under.
onnx_model_filename = "resnet18-4w4a_tidied"

# Wrap our onnx model in the ModelWrapper class
onnx_model = ModelWrapper(f"{onnx_model_filename}.onnx")

# Use the ModelWrapper class' transform function to
# apply the MoveLinearPastEltWiseAdd() transformation.
onnx_model = onnx_model.transform(MoveLinearPastEltwiseAdd())

# Save the modified ONNX model.
onnx_model.save(f"{onnx_model_filename}_MODIFIED.onnx")

5. Start the docker container and have it run 'build.py', with the command: ./run-docker.sh build_custom <Path_To_Build.py>. Make sure to replace <Path_To_Build.py> with the directory containing 'build.py' and 'resnet18-4w4a_tidied.onnx'.
6. Wait for the docker container to finish running 'build.py'. It will print The program finished and will be restarted once it has completed.
7. Open the modified ONNX model in Netron.app

This process produces the following ONNX file: 'resnet18-4w4a_tidied_MODIFIED.onnx'

Expected behavior

Either:

• Before moving linear nodes past element-wise addition nodes, The transformation moves all forked linear operations past their forks, with their producers becoming forks instead (this is already implemented in MoveLinearPastFork()). This would avoid the transformation moving fork nodes.
  or...
• The transformation warns the user if they try to apply MoveLinearPastEltWiseAdd() in a way that would result in fork nodes being moved past join nodes.
  or...
• The transformation ignores linear fork nodes.

Actual behavior

The transformation moves a fork node past its relevant joint node. This results in a huge mess of a graph (see images below).

• Before applying MoveLinearPastEltWiseAdd():
  
• After applying MoveLinearPastEltWiseAdd():
  

Possible fixes

I'm unsure which of these three fixes would suit the project best:

• MoveLinearPastFork() could be applied to the model at the start of the transformation. This moves all forked linear operations past their forks, with their producers becoming forks instead.
• A warning could be raised whenever a fork node is moved by the transformation.
• The transformation could ignore fork nodes.

Additional context

• The above behaviour is technically correct!
  • In the example there are 2 identical multiply nodes before an element-wise addition. The transformation's correct behaviour is to move both of them past the element-wise add. However, the transformation does not realise one of the multiply nodes is a fork. The graph produced in this instance is complete nonsense.
• While the behavior is "correct", (and easily circumvented by I don't think it's expected.

ONNX files

Pre-transformation ONNX file: 'resnet18-4w4a_tidied.onnx'
Post-transformation ONNX file: 'resnet18-4w4a_tidied_MODIFIED.onnx'