Seemingly incorrect canonicalization for SwitchOp

MLIR
1

The following canonicalized with mlir-opt -canonicalize:

func @switch(%flag : i32, %caseOperand : i32, %t1 : tensor<f32>, %t2 : tensor<f32>)
    -> (tensor<f32>) 
{
  switch %flag : i32, [
    default: ^bb1(%caseOperand : i32),
    42: ^bb2(%caseOperand : i32)
  ]

  ^bb1(%bb1arg : i32):
    return %t1 : tensor<f32>
  ^bb2(%bb2arg : i32):
    return %t2 : tensor<f32>
}

produces:

func @switch(%arg0: i32, %arg1: i32, %arg2: tensor<f32>, %arg3: tensor<f32>) -> tensor<f32> {
  switch %arg0 : i32, [
    default: ^bb1(%arg2 : tensor<f32>),
    42: ^bb1(%arg3 : tensor<f32>)
  ]
 ^bb1(%0: tensor<f32>):  // 2 preds: ^bb0, ^bb0
  return %0 : tensor<f32>
}

Do folks see a mistake in my example or is this a bug?

2

What do you believe is incorrect here? Switch is on first argument, default case returns third argument, 42 returns fourth in both before and after, while 2nd argument unused in after.

3

you’re right I misread the change of type in the bbargs, the folding is more powerful than I expected :slight_smile: