[MLIR] Can't lower tosa.while_loop

MLIR TOSA
1

Hello!

// from tosa-to-scf test
func.func @while_test(%arg0 : tensor<i32>) -> (tensor<i32>) {
  %0 = tosa.while_loop (%arg1 = %arg0) : (tensor<i32>) -> tensor<i32> {
    %1 = "tosa.const"() {value = dense<3> : tensor<i32>} : () -> tensor<i32>
    %2 = tosa.greater_equal %1, %arg1 : (tensor<i32>, tensor<i32>) -> tensor<i1>
    tosa.yield %2 : tensor<i1>
  } do {
  ^bb0(%arg1: tensor<i32>):
    %1 = "tosa.const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
    %2 = tosa.add %arg1, %1 : (tensor<i32>, tensor<i32>) -> tensor<i32>
    tosa.yield %2 : tensor<i32>
  }
  return %0 : tensor<i32>
}

When I lower this IR by mlir-opt -–tosa-to-scf --one-shot-bufferize test.mlir, I get the following error:

error: Yield operand #0 is not equivalent to the corresponding iter bbArg

Am I doing something wrong here? The (now deleted) scf-bufferize pass seems to work fine

2

P.S. the problem is resolved, I just needed to set allowReturnAllocsFromLoops to true in the pass options.