[regression] error: failed to legalize operation 'arith.sitofp' that was explicitly marked illegal

[pdhirajkumarprasad]

This model was working till 1114 and started failing from 1115, seems like torch-mlir bump has caused this issue

module {

  func.func @main_graph(%arg0: !torch.vtensor<[?,84,?,?],f32> ) -> !torch.vtensor<[],f32>  attributes {torch.onnx_meta.ir_version = 8 : si64, torch.onnx_meta.opset_version = 21 : si64, torch.onnx_meta.producer_name = "pytorch", torch.onnx_meta.producer_version = "2.1.0"} {

    %none = torch.constant.none

    %1 = torch.operator "onnx.Relu"(%arg0) : (!torch.vtensor<[?,84,?,?],f32>) -> !torch.vtensor<[?,84,?,?],f32> 

    %2 = torch.operator "onnx.Shape"(%1) : (!torch.vtensor<[?,84,?,?],f32>) -> !torch.vtensor<[4],si64> 

    %3 = torch.operator "onnx.Constant"() {torch.onnx.value = dense_resource<__186> : tensor<si64>} : () -> !torch.vtensor<[],si64> 

    %4 = torch.operator "onnx.Gather"(%2, %3) {torch.onnx.axis = 0 : si64} : (!torch.vtensor<[4],si64>, !torch.vtensor<[],si64>) -> !torch.vtensor<[],si64> 

    %5 = torch.operator "onnx.Shape"(%1) : (!torch.vtensor<[?,84,?,?],f32>) -> !torch.vtensor<[4],si64> 

    %6 = torch.operator "onnx.Constant"() {torch.onnx.value = dense_resource<__187> : tensor<si64>} : () -> !torch.vtensor<[],si64> 

    %8 = torch.operator "onnx.Cast"(%4) {torch.onnx.to = 1 : si64} : (!torch.vtensor<[],si64>) -> !torch.vtensor<[],f32> 

    %9 = torch.operator "onnx.Constant"() {torch.onnx.value = dense_resource<__206> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64> 

    %10 = torch.operator "onnx.Unsqueeze"(%8, %9) : (!torch.vtensor<[],f32>, !torch.vtensor<[1],si64>) -> !torch.vtensor<[1],f32> 

    %11 = torch.operator "onnx.Concat"(%10, %10, %10, %10) {torch.onnx.axis = 0 : si64} : (!torch.vtensor<[1],f32>, !torch.vtensor<[1],f32>, !torch.vtensor<[1],f32>, !torch.vtensor<[1],f32>) -> !torch.vtensor<[4],f32> 

    %12 = torch.operator "onnx.Cast"(%11) {torch.onnx.to = 7 : si64} : (!torch.vtensor<[4],f32>) -> !torch.vtensor<[4],si64> 

    %13 = torch.operator "onnx.Concat"(%12, %12) {torch.onnx.axis = 0 : si64} : (!torch.vtensor<[4],si64>, !torch.vtensor<[4],si64>) -> !torch.vtensor<[8],si64> 

    %14 = torch.operator "onnx.Cast"(%13) {torch.onnx.to = 7 : si64} : (!torch.vtensor<[8],si64>) -> !torch.vtensor<[8],si64> 

    %15 = torch.operator "onnx.Constant"() {torch.onnx.value = dense_resource<__219> : tensor<f32>} : () -> !torch.vtensor<[],f32> 

    %16 = torch.operator "onnx.Pad"(%1, %14, %15) {torch.onnx.mode = "constant"} : (!torch.vtensor<[?,84,?,?],f32>, !torch.vtensor<[8],si64>, !torch.vtensor<[],f32>) -> !torch.vtensor<[?,?,?,?],f32> 

    return %8 : !torch.vtensor<[],f32>

  }

}



{-#

  dialect_resources: {

    builtin: {

      __186: "0x080000000200000000000000",

      __187: "0x080000000300000000000000",

      __206: "0x080000000000000000000000",

      __219: "0x0800000000000000"

    }

  }

#-}

[vivekkhandelwal1]

The issue is because of the following changes: llvm/torch-mlir@cd38ecf#diff-7b9f4faa909a47945ec2c3e11de440fa77512517353d4fb922fc4efa69a06480R714-R851

Working IR (without these changes):

#map = affine_map<() -> ()>
module {
  ml_program.global private mutable @global_seed(dense<0> : tensor<i64>) : tensor<i64>
  func.func @main_graph(%arg0: tensor<?x84x?x?xf32>) -> tensor<f32> {
    %c2 = arith.constant 2 : index
    %dim = tensor.dim %arg0, %c2 : tensor<?x84x?x?xf32>
    %0 = arith.index_cast %dim : index to i64
    %1 = tensor.empty() : tensor<i64>
    %2 = linalg.fill ins(%0 : i64) outs(%1 : tensor<i64>) -> tensor<i64>
    %3 = tensor.empty() : tensor<f32>
    %4 = linalg.generic {indexing_maps = [#map, #map], iterator_types = []} ins(%2 : tensor<i64>) outs(%3 : tensor<f32>) {
    ^bb0(%in: i64, %out: f32):
      %5 = arith.sitofp %in : i64 to f32
      linalg.yield %5 : f32
    } -> tensor<f32>
    return %4 : tensor<f32>
  }
}

Non-working IR (with these changes):

module {
  ml_program.global private mutable @global_seed(dense<0> : tensor<i64>) : tensor<i64>
  func.func @main_graph(%arg0: tensor<?x84x?x?xf32>) -> tensor<f32> {
    %c2 = arith.constant 2 : index
    %dim = tensor.dim %arg0, %c2 : tensor<?x84x?x?xf32>
    %0 = arith.index_cast %dim : index to i64
    %1 = arith.sitofp %0 : i64 to f64
    %2 = arith.truncf %1 : f64 to f32
    %3 = tensor.empty() : tensor<f32>
    %4 = linalg.fill ins(%2 : f32) outs(%3 : tensor<f32>) -> tensor<f32>
    return %4 : tensor<f32>
  }
}

The particular thing to notice is that the way i64 -> f32 conversion is happening, that makes a difference and results in this failure.

[pdhirajkumarprasad]

Following 2 models are impacted

migraphx_sd__unet__model
migraphx_sdxl__unet__model

[pashu123]

@pdhirajkumarprasad It works after adding -iree-input-demote-i64-to-i32 this pass to the compile command. Also, it's a fair assumption since the index comes from tensor.dim. I am currently looking into adding support for arith.sitofp %0 : i64 to f32 ops in the VM conversion.

[pashu123]

@vinayakdsci is adding support for the new vm op that's missing and its c dispatch function.

[vinayakdsci]

Resolved by iree-org/iree#19455.

[zjgarvey]

Let's check the report for those models. If they are failing due a different issue, then we can close this.