TOSA nested while loop

MLIR
1

Hi all,

Can anyone help me with the following tosa dialect IR? I feel it should be able to convert via mlir-opt --tosa-to-scf but I get a failure to legalize operation when I run it through mlir-op.

module  {                                                                                                                              
  func @entry_fun(%ival: tensor<i32>, %jval: tensor<i32>, %lim: tensor<i32>) -> tensor<i32> {                                          
    %0 = "tosa.while_loop"(%ival) ( {                                                                                                  
    ^bb0(%arg0: tensor<i32>):  // no predecessors                                                                                      
      %1 = "tosa.greater_equal"(%arg0, %lim) : (tensor<i32>, tensor<i32>) -> tensor<i1>                                                
      %2 = "tosa.bitwise_not"(%1) : (tensor<i1>) -> tensor<i1>                                                                         
      "tosa.yield"(%2, %arg0) : (tensor<i1>, tensor<i32>) -> ()                                                                        
    },  {                                                                                                                              
    ^bb0(%arg0: tensor<i32>):  // no predecessors                                                                                      
      %2 = "tosa.while_loop"(%jval) ( {                                                                                                
      ^bb0(%arg1: tensor<i32>):  // no predecessors                                                                                    
        %3 = "tosa.greater_equal"(%arg1, %lim) : (tensor<i32>, tensor<i32>) -> tensor<i1>                                              
        %4 = "tosa.bitwise_not"(%3) : (tensor<i1>) -> tensor<i1>                                                                       
        "tosa.yield"(%4, %arg1) : (tensor<i1>, tensor<i32>) -> ()                                                                      
      },  {                                                                                                                            
      ^bb0(%arg1: tensor<i32>):  // no predecessors                                                                                    
        %3 = "tosa.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>                                                      
        "tosa.yield"(%3) : (tensor<i32>) -> ()                                                                                         
      }) : (tensor<i32>) -> tensor<i32>                                                                                                
      "tosa.yield"(%2) : (tensor<i32>) -> ()                                                                                           
    }) : (tensor<i32>) -> tensor<i32>                                                                                                  
    return %0 : tensor<i32>                                                                                                            
  }                                                                                                                                    
}

Is it possible to have nested tosa.while_loops ops? Unless I am missing something very obvious here I am a little stumped as to why the IR is not converting.

Thanks!
Anthony

2

Hey Anthony,

What is the error that you see? Perhaps would be good to make a github issue for it too (I think it is allowed and initial look at your IR, it seemed reasonable).

– Jacques

3

Hi Jacques,

I get the following error:

test/simp.mlir:3:10: error: failed to legalize operation 'tosa.while_loop' that was explicitly marked illegal
    %0 = "tosa.while_loop"(%ival) ( {
         ^
test/simp.mlir:3:10: note: see current operation: %0 = "tosa.while_loop"(%arg0) ( {
^bb0(%arg3: tensor<i32>):  // no predecessors
  %1 = "tosa.greater_equal"(%arg3, %arg2) : (tensor<i32>, tensor<i32>) -> tensor<i1>
  %2 = "tosa.bitwise_not"(%1) : (tensor<i1>) -> tensor<i1>
  "tosa.yield"(%2, %arg3) : (tensor<i1>, tensor<i32>) -> ()
},  {
^bb0(%arg3: tensor<i32>):  // no predecessors
  %1 = "tosa.while_loop"(%arg1) ( {
  ^bb0(%arg4: tensor<i32>):  // no predecessors
    %2 = "tosa.greater_equal"(%arg4, %arg2) : (tensor<i32>, tensor<i32>) -> tensor<i1>
    %3 = "tosa.bitwise_not"(%2) : (tensor<i1>) -> tensor<i1>
    "tosa.yield"(%3, %arg4) : (tensor<i1>, tensor<i32>) -> ()
  },  {
  ^bb0(%arg4: tensor<i32>):  // no predecessors
    %2 = "tosa.add"(%arg3, %arg4) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    "tosa.yield"(%2) : (tensor<i32>) -> ()
  }) : (tensor<i32>) -> tensor<i32>
  "tosa.yield"(%1) : (tensor<i32>) -> ()
}) : (tensor<i32>) -> tensor<i32>

Which is why I am a little confused, because it looks like this pops up during the conversion to SCF. If I run with -debug, I get…

    //===-------------------------------------------===//
    Legalizing operation : 'tosa.while_loop'(0x314a150) {
      * Fold {
      } -> FAILURE : unable to fold

      * Pattern : 'tosa.while_loop -> ()' {
      } -> FAILURE : pattern was already applied
    } -> FAILURE : no matched legalization pattern
    //===-------------------------------------------===//
  } -> FAILURE : generated operation 'tosa.while_loop'(0x000000000314A150) was illegal
} -> FAILURE : pattern failed to match
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

Does it still seem like a bug? If so I’m happy to open an issue.

4

I took a brief glance at the code and looks like this is a combination of: the lowering for tosa.while clones the region (not sure why this isn’t just inlining it), and because the lowering for tosa.while does not indicate that it can be applied recursively (Pattern Rewriting : Generic DAG-to-DAG Rewriting - MLIR, e.g. llvm-project/VectorToSCF.cpp at 6f1ce046fdd02458fe2e3824b865ba4bdf129d67 · llvm/llvm-project · GitHub). Given that it clones the region, all operations that were cloned are immediately considered for legalization. This then fails because we try to apply the same tosa.while pattern, which doesn’t support recursion. Fixing either of the two things (not cloning, or allowing recursion) above would likely fix the issue.

– River

5

A side note regarding this. I recently hit this same issue.

  1. I chose not to inline because inlining a region results in the old op becoming malformed, and -debug will result in segfaults without remembering to specify -mlir-print-op-generic.
  2. Cloning an operation’s body makes visible this bug, that nested ops of the same type cannot be differentiated from a recursive application of a pattern.

Just pointing that out in case it changes any priorities and explaining why one might clone rather than inline.

6

This is an infra problem to solve, using clone does not seems like a good solution: it is much more expensive to do.
Did you file a bug for this?

7

I have now filed one here: inlineRegionBefore breaks -debug · Issue #53921 · llvm/llvm-project · GitHub

1 Like