onnx-mlir
onnx-mlir copied to clipboard
onnx
If
Implemented verify() and inferShapes() for If op. Tested with lit tests.
Also reordered the else and then branches so that
%0 = "onnx.If"(%arg0) ({
%1 = "onnx.Constant"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
onnx.Return %1, %2 : tensor<i32>
}, {
%1 = "onnx.Constant"() {value = dense<2> : tensor<i32>} : () -> tensor<i32>
onnx.Return %1 : tensor<i32>
}) : (tensor<i1>) -> tensor<i32>
has the then branch before the else branch and not the other way around, which I found very counterintuitive. Curiously, the ONNX spec lists the else_branch before the then_branch attribute, so it wasn't enough to not sort the attributes alphabetically and I put in special logic in gen_onnx_mlir.py to put then_branch before else_branch, controlled by the special_attr_order set.
@chentong319 added you as reviewer since I believe if is similar to loop, from what I recall from Tian's conversations.
@chentong319 do you mind looking into this one, as it is related to loop? tx
To keep the code simple, I prefer not to change the index for the 'then' and 'else' subgraph. To me, the index is anyway implementation dependent and the string name for the subgraph is always used. The original order is simply lexicographic order of the name.
To keep the code simple, I prefer not to change the index for the 'then' and 'else' subgraph. To me, the index is anyway implementation dependent and the string name for the subgraph is always used. The original order is simply lexicographic order of the name.
I really think that people are likely to read "onnx.If(%1) ({A}, {B}) to mean if %1 then A else B and therefore it is too confusing if we keep it the other way around. I checked the 'if' constructs in MLIR (scf.if, affine.if, tosa.cond_if) and they all have 'then' before 'else'.
I appreciate that you want to keep the code simple so I refactored the code in gen_onnx_mlir.py to stash the complexity into a standalone funtion order_attr_names and I polished the implementation and added comments to make it easier to understand. PTAL.
Of course, if I can't change your mind just let me know and I will remove it.
I agree that the output onnx.If(%1) ({A}, {B}) is confusing to user. Previously I thought the index is only used in accessing the then/else region in code. It may require more change if we customize the printer and parser for If Op.
