Structured conditional construct
ToucheSir opened this issue · comments
Being able to represent https://github.com/onnx/onnx/blob/master/docs/Operators.md#if would be great. I think it might also remove some of the sharper edges of normal tracing, such as certain branches not being hit.
Supporting If operation on Tape is relatively easy - we already have a more complicated Loop anyway. However, supporting the branches that have not been hit during the tracing is a way more complicated. The tracer essentially works by adding callbacks to different parts of the IR code and then constructing a tape in these callbacks. If some piece of code isn't hit, the corresponding callback isn't invoked and the piece of code doesn't get to the tape. For the same reason Loop only works correctly if its body is executed at least once during tracing (see the warning at the end of Loop docs).
Instead of tracing, i.e. recording operations during execution, we could go the Mjolnir way and translate IR into a tape fully statically. I have certain plans to try this out, but given that IRTools.jl is effectively in maintenance mode and CompilerPlugins are still in development, I don't rush with this activity. Fortunately, dynamic computational graphs are rare enough in practice that we can postpone them for quite some time.
For sure, I wanted to put a pin in this for if/when a more mature compiler plugin stack emerges. Most users are likely used to the control flow limitation from PyTorch's non-jit ONNX export anyhow.