Issues with backend test: test_loop13_seq

[chentong319]

Test_loop13_seq.onnx can not be compiled. The main reason is !onnx.Seq<tensor<*xf32>>. I just record my investigation with this issue. I do not think that we will work on this issue in near future.

The EmitONNXBasic output of model test_loop13_seq.onnx:

  func @main_graph(%arg0: tensor<i64>, %arg1: tensor<i1>, %arg2: !onnx.Seq<tensor<f32>>) -> tensor<i1> attributes {input_names = ["trip_count", "cond", "seq_empty"], output_names = ["seq_res"]} {
    %0 = "onnx.Loop"(%arg0, %arg1, %arg2) ({
    ^bb0(%arg3: tensor<i64>, %arg4: tensor<i1>, %arg5: !onnx.Seq<tensor<*xf32>>):
      %1 = "onnx.Identity"(%arg4) : (tensor<i1>) -> tensor<i1>
      %2 = "onnx.Constant"() {value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00]> : tensor<5xf32>} : () -> tensor<5xf32>
      %3 = "onnx.Constant"() {value = dense<1> : tensor<i64>} : () -> tensor<i64>
      %4 = "onnx.Constant"() {value = dense<0> : tensor<1xi64>} : () -> tensor<1xi64>
      %5 = "onnx.Add"(%arg3, %3) : (tensor<i64>, tensor<i64>) -> tensor<*xi64>
      %6 = "onnx.Constant"() {value = dense<0> : tensor<i64>} : () -> tensor<i64>
      %7 = "onnx.Unsqueeze"(%5, %6) : (tensor<*xi64>, tensor<i64>) -> tensor<*xi64>
      %8 = "onnx.NoValue"() {value} : () -> none
      %9 = "onnx.NoValue"() {value} : () -> none
      %10 = "onnx.Slice"(%2, %4, %7, %8, %9) : (tensor<5xf32>, tensor<1xi64>, tensor<*xi64>, none, none) -> tensor<*xf32>
      %11 = "onnx.NoValue"() {value} : () -> none
      %12 = "onnx.SequenceInsert"(%arg5, %10, %11) : (!onnx.Seq<tensor<*xf32>>, tensor<*xf32>, none) -> !onnx.Seq<tensor<*xf32>>
      onnx.Return %1, %12 : tensor<i1>, !onnx.Seq<tensor<*xf32>>
    }) {input_names = ["iter_count", "cond_in", "seq_in"], output_names = ["cond_out", "seq_out"]} : (tensor<i64>, tensor<i1>, !onnx.Seq<tensor<f32>>) -> tensor<i1>
    return %0 : tensor<i1>
  }

The unranked tensor, often used as sequence element type for a generic sequence, can not be handled by onnx-mlir currently. Not only there is a check for unranked tensor at the end of ONNX phase, but also the krnl allows only memref, not unranked memref, in many of its ops and utilities. For example, mem.cast requires the type pavement to be memref. If that restriction is lifted, I ran into issue in krnl.store. A sequence of unranked tensor itself could be easily handled in sequence operations, because it is just a pointer. However, the tensor retrieved from a such sequence will be an unranked tensor that onnx-mlir can not handle. Thus, I will not pursue sequence of unranked tensor now. In the model test_loop13_seq, the input sequence (%arg2) has type Seq<tensor> , and tensors with type tensor<*xf32> is inserted into the sequence in onnx.loop. Even though shape inference can find out that the inserted tensor is of type tensor<?xf32> , the sequence has to be Seq<tensor<*xf32>>. If the input sequence (%arg2) is modified as seq<tensor<?xf32>>, the shape inference is able to infer the sequence to be seq<tensor<?xf32>>. Modified mlir:

  func @main_graph(%arg0: tensor<i64>, %arg1: tensor<i1>, %arg2: !onnx.Seq<tensor<?xf32>>) -> tensor<i1> attributes {input_names = ["trip_count", "cond", "seq_empty"], output_names = ["seq_res"]} {
    %0 = "onnx.Loop"(%arg0, %arg1, %arg2) ({
    ^bb0(%arg3: tensor<i64>, %arg4: tensor<i1>, %arg5: !onnx.Seq<tensor<*xf32>>):
      %1 = "onnx.Identity"(%arg4) : (tensor<i1>) -> tensor<i1>
      %2 = "onnx.Constant"() {value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00]> : tensor<5xf32>} : () -> tensor<5xf32>
      %3 = "onnx.Constant"() {value = dense<1> : tensor<i64>} : () -> tensor<i64>
      %4 = "onnx.Constant"() {value = dense<0> : tensor<1xi64>} : () -> tensor<1xi64>
      %5 = "onnx.Add"(%arg3, %3) : (tensor<i64>, tensor<i64>) -> tensor<*xi64>
      %6 = "onnx.Constant"() {value = dense<0> : tensor<i64>} : () -> tensor<i64>
      %7 = "onnx.Unsqueeze"(%5, %6) : (tensor<*xi64>, tensor<i64>) -> tensor<*xi64>
      %8 = "onnx.NoValue"() {value} : () -> none
      %9 = "onnx.NoValue"() {value} : () -> none
      %10 = "onnx.Slice"(%2, %4, %7, %8, %9) : (tensor<5xf32>, tensor<1xi64>, tensor<*xi64>, none, none) -> tensor<*xf32>
      %11 = "onnx.NoValue"() {value} : () -> none
      %12 = "onnx.SequenceInsert"(%arg5, %10, %11) : (!onnx.Seq<tensor<*xf32>>, tensor<*xf32>, none) -> !onnx.Seq<tensor<*xf32>>
      onnx.Return %1, %12 : tensor<i1>, !onnx.Seq<tensor<*xf32>>
    }) {input_names = ["iter_count", "cond_in", "seq_in"], output_names = ["cond_out", "seq_out"]} : (tensor<i64>, tensor<i1>, !onnx.Seq<tensor<?xf32>>) -> tensor<i1>
    return %0 : tensor<i1>
  }

Output after shape inference:

  func @main_graph(%arg0: tensor<i64>, %arg1: tensor<i1>, %arg2: !onnx.Seq<tensor<?xf32>>) -> !onnx.Seq<tensor<?xf32>> attributes {input_names = ["trip_count", "cond", "seq_empty"], output_names = ["seq_res"]} {
    %0 = "onnx.Loop"(%arg0, %arg1, %arg2) ({
    ^bb0(%arg3: tensor<i64>, %arg4: tensor<i1>, %arg5: !onnx.Seq<tensor<?xf32>>):
      %1 = "onnx.Identity"(%arg4) : (tensor<i1>) -> tensor<i1>
      %2 = "onnx.Constant"() {value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00]> : tensor<5xf32>} : () -> tensor<5xf32>
      %3 = "onnx.Constant"() {value = dense<1> : tensor<i64>} : () -> tensor<i64>
      %4 = "onnx.Constant"() {value = dense<0> : tensor<1xi64>} : () -> tensor<1xi64>
      %5 = "onnx.Add"(%arg3, %3) : (tensor<i64>, tensor<i64>) -> tensor<i64>
      %6 = "onnx.Constant"() {value = dense<0> : tensor<i64>} : () -> tensor<i64>
      %7 = "onnx.Unsqueeze"(%5, %6) : (tensor<i64>, tensor<i64>) -> tensor<1xi64>
      %8 = "onnx.NoValue"() {value} : () -> none
      %9 = "onnx.NoValue"() {value} : () -> none
      %10 = "onnx.Constant"() {value = dense<0> : tensor<1xi64>} : () -> tensor<1xi64>
      %11 = "onnx.Constant"() {value = dense<1> : tensor<1xi64>} : () -> tensor<1xi64>
      %12 = "onnx.Slice"(%2, %4, %7, %10, %11) : (tensor<5xf32>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> tensor<?xf32>
      %13 = "onnx.NoValue"() {value} : () -> none
      %14 = "onnx.SequenceInsert"(%arg5, %12, %13) : (!onnx.Seq<tensor<?xf32>>, tensor<?xf32>, none) -> !onnx.Seq<tensor<?xf32>>
      onnx.Return %1, %14 : tensor<i1>, !onnx.Seq<tensor<?xf32>>
    }) {input_names = ["iter_count", "cond_in", "seq_in"], output_names = ["cond_out", "seq_out"]} : (tensor<i64>, tensor<i1>, !onnx.Seq<tensor<?xf32>>) -> !onnx.Seq<tensor<?xf32>>
    return %0 : !onnx.Seq<tensor<?xf32>>
  }

With the modification, there is no seq<tensor<*xf32>> at lowering from onnx to krnl. However, an error occurred for return: the Seq type is not lowered. No problem if the model returns a tensor retrieved from the seq. Further investigation needed for the return. The output of the return is here:

    %14 = "builtin.unrealized_conversion_cast"(%3) : (memref<?xmemref<?xf32>>) -> !onnx.Seq<tensor<?xf32>>
    "func.return"(%14) : (!onnx.Seq<tensor<?xf32>>) -> ()