error: 'func.func' op unhandled function with multiple blocks

[pashu123]

What happened?

latest.mlir:3:10: error: 'func.func' op unhandled function with multiple blocks
    %0 = flow.dispatch.region -> (tensor<2x4xf32>) {
         ^
latest.mlir:2:3: note: called from
  func.func @simple_test_with_cfg(%arg0: i8) -> tensor<2x4xf32> {
  ^
latest.mlir:3:10: note: see current operation: 
"func.func"() <{function_type = () -> (), sym_name = "simple_test_with_cfg_dispatch_0"}> ({
  %0 = "arith.constant"() <{value = 0 : index}> : () -> index
  %1 = "arith.constant"() <{value = 0 : i8}> : () -> i8
  %2 = "arith.constant"() <{value = dense<1.000000e+00> : tensor<2x4xf32>}> : () -> tensor<2x4xf32>
  %3 = "bufferization.to_memref"(%2) : (tensor<2x4xf32>) -> memref<2x4xf32>

https://gist.github.com/pashu123/4cf00bdd0eea5f76043c6b64e9ed5d08

Steps to reproduce your issue

Branch: https://github.com/pashu123/iree/tree/forallmmt

Example IR:

module {
  func.func @simple_test_with_cfg(%arg0: i8) -> tensor<2x4xf32> {
    %0 = flow.dispatch.region -> (tensor<2x4xf32>) {
      %c0_i8 = arith.constant 0 : i8
      %1 = arith.cmpi eq, %arg0, %c0_i8 : i8
      cf.cond_br %1, ^bb1, ^bb2
    ^bb1:  // pred: ^bb0
      %2 = tensor.empty() : tensor<2x4xf32>
      flow.return %2 : tensor<2x4xf32>
    ^bb2:  // pred: ^bb0
      %cst = arith.constant dense<1.000000e+00> : tensor<2x4xf32>
      flow.return %cst : tensor<2x4xf32>
    }
    return %0 : tensor<2x4xf32>
  }
}

Command: iree-compile --iree-hal-target-backends=llvm-cpu --iree-llvmcpu-target-cpu=znver4 latest.mlir -o out_new.vmfb

What component(s) does this issue relate to?

No response

Version information

No response

Additional context

No response