Transform SequenceAt to split for special cases

[chentong319]

Certain PyTorch.onnx.export will break the LSTM op into lower level operations and generate SplitToSequence and SequenceAt operation pattern. For example:

  %15 = onnx.Constant dense<0> : tensor<i64>
  %38 = onnx.Constant dense<1> : tensor<i64>
  %65 = onnx.Constant dense<100> : tensor<i64>
  %66 = "onnx.SplitToSequence"(%arg0, %65) {axis = 2 : si64, keepdims = 1 : si64} : (tensor<1x1x400xf32>, tensor<i64>) -> !onnx.Seq<tensor<1x1x100xf32>>
  %67 = "onnx.SequenceAt"(%66, %15) : (!onnx.Seq<tensor<1x1x100xf32>>, tensor<i64>) -> tensor<1x1x100xf32>
  %68 = "onnx.SequenceAt"(%66, %38) : (!onnx.Seq<tensor<1x1x100xf32>>, tensor<i64>) -> tensor<1x1x100xf32>
  %40 = "onnx.Add"(%67, %68) : (tensor<1x1x100xf32>, tensor<1x1x100xf32>) -> tensor<1x1x100xf32>

ONNX-MLIR currently does not have the lowering for SplitToSequence. In general, sequence related ops are not well optimized in ONNX-MLIR. I found that such code pattern can be converted into tensor operations to avoid the sequence ops.

    %0 = onnx.Constant dense<100> : tensor<4xi64>
    %1:4 = "onnx.Split"(%arg0, %0) {axis = 2 : si64} : (tensor<1x1x400xf32>, tensor<4xi64>) -> (tensor<1x1x100xf32>, tensor<1x1x100xf32>, tensor<1x1x100xf32>, tensor<1x1x100xf32>)
    %2:4 = "onnx.Split"(%arg0, %0) {axis = 2 : si64} : (tensor<1x1x400xf32>, tensor<4xi64>) -> (tensor<1x1x100xf32>, tensor<1x1x100xf32>, tensor<1x1x100xf32>, tensor<1x1x100xf32>)
    %3 = "onnx.Add"(%1#0, %2#1) : (tensor<1x1x100xf32>, tensor<1x1x100xf32>) -> tensor<1x1x100xf32>

The two onnx.Split is supposed to be able merged into one. But current onnx-mlir didn't. Need further investigation.

But the exported model has another type of SequenceAt, in which the output type of SequenceAt is different from the element type of the input sequence type. I think it is an error in the exporter. However, I tried to fix this issue in the transformation of this PR.

    %26 = onnx.Constant dense<0> : tensor<i64>
    %27 = onnx.Constant dense<1> : tensor<i64>
    %32 = "onnx.SplitToSequence"(%arg0, %27) {axis = 0 : si64, keepdims = 0 : si64} : (tensor<1x1x100xf32>, tensor<i64>) -> !onnx.Seq<tensor<1x1x100xf32>>
    %33 = "onnx.SequenceAt"(%32, %26) : (!onnx.Seq<tensor<1x1x100xf32>>, tensor<i64>) -> tensor<1x100xf32>

Output from this PR:

    %0 = onnx.Constant dense<0> : tensor<1xi64>
    %1 = onnx.Constant dense<1> : tensor<1xi64>
    %2 = "onnx.Split"(%arg0, %1) {axis = 0 : si64} : (tensor<1x1x100xf32>, tensor<1xi64>) -> tensor<1x1x100xf32>
    %3 = "onnx.Squeeze"(%2, %0) : (tensor<1x1x100xf32>, tensor<1xi64>) -> tensor<1x100xf32>

Other small fixes in this PR:

1. Handle the case that onnx.ConstantOp has not been normalized yet. Otherwise, --EmitONNXBasic will fail for this model.
2. It is not needed to check the tensor range size to be larger than one. In the model, the number of output from Split happened to be one.

Test:
With same random input, the outputs of lstm_no_data and lstm_no_dynamo are the same for atol=0.01, rtol=0.05, with or without NNPA turned on.

[tungld]

LGTM.

[tungld]

Can be shorten with

ONNXSplitToSequenceOp splitToSequence = inputSequence.getDefiningOp<ONNXSplitToSequenceOp>();
if (!splitToSequence)
  return false;

[chentong319]

Done

[tungld]

Same here, we can use ONNXSequenceAtOp op = sequenceAtResult.getDefiningOp<ONNXSequenceAtOp>()

[chentong319]

Done

[tungld]

Look like we can remove this completely.

[chentong319]

Done

[AlexandreEichenberger]

LGTM, thanks for the quick fix and the lit-test.

It would be interesting to study how this pattern fairs compared to the original pattern (it was for an LSTM, correct?)

[AlexandreEichenberger]

@chentong319 Dominic is eager to test this. Can you implement the minor suggestions and merge it into main? Thanks

[AlexandreEichenberger]

thanks @chentong319 .

[jenkins-droid]

Jenkins Linux amd64 Build #16051 [push] Transform SequenceAt to ... started at 12:34

[jenkins-droid]

Jenkins Linux s390x Build #16054 [push] Transform SequenceAt to ... started at 13:34

[jenkins-droid]

Jenkins Linux ppc64le Build #15081 [push] Transform SequenceAt to ... started at 13:50

[jenkins-droid]

Jenkins Linux amd64 Build #16051 [push] Transform SequenceAt to ... failed after 1 min 36 sec

[jenkins-droid]

Jenkins Linux s390x Build #16054 [push] Transform SequenceAt to ... passed after 1 hr 41 min

[jenkins-droid]

Jenkins Linux ppc64le Build #15081 [push] Transform SequenceAt to ... passed after 2 hr 18 min