Failed to legalize operation 'torch.operator' that was explicitly marked illegal : "onnx.Loop"

[pravg-amd]

For the below IR,

func.func @main(%arg0: !torch.vtensor<[],si64>, %arg1: !torch.vtensor<[],i1>, %arg2: !torch.vtensor<[],si32>, %arg3: !torch.vtensor<[],si32>) -> (!torch.vtensor<[],si32>, !torch.vtensor<[?],si32>) attributes {torch.onnx_meta.ir_version = 6 : si64, torch.onnx_meta.opset_version = 21 : si64, torch.onnx_meta.producer_name = "keras2onnx", torch.onnx_meta.producer_version = "1.6.0"} {
  %0:2 = torch.operator "onnx.Loop"(%arg0, %arg1, %arg2) : (!torch.vtensor<[],si64>, !torch.vtensor<[],i1>, !torch.vtensor<[],si32>) -> (!torch.vtensor<[],si32>, !torch.vtensor<[?],si32>) {
    ^bb0(%arg4: !torch.vtensor<[],si64>, %arg5: !torch.vtensor<[],i1>, %arg6: !torch.vtensor<[],si32>):
      %592 = torch.operator "onnx.Add"(%arg6, %arg3) : (!torch.vtensor<[],si32>, !torch.vtensor<[],si32>) -> !torch.vtensor<[],si32>
      %593 = torch.operator "onnx.Identity"(%arg6) : (!torch.vtensor<[],si32>) -> !torch.vtensor<[],si32>
      %594 = torch.operator "onnx.Identity"(%arg5) : (!torch.vtensor<[],i1>) -> !torch.vtensor<[],i1>
      torch.operator_terminator %594, %592, %593 : !torch.vtensor<[],i1>, !torch.vtensor<[],si32>, !torch.vtensor<[],si32>
  }
  return %0#0, %0#1 : !torch.vtensor<[],si32>, !torch.vtensor<[?],si32>
}

Getting error as

model_torch_onnx.mlir:2:10: error: failed to legalize operation 'torch.operator' that was explicitly marked illegal
  %0:2 = torch.operator "onnx.Loop"(%arg0, %arg1, %arg2) : (!torch.vtensor<[],si64>, !torch.vtensor<[],i1>, !torch.vtensor<[],si32>) -> (!torch.vtensor<[],si32>, !torch.vtensor<[?],si32>) {
         ^
model_torch_onnx:2:10: note: see current operation:
%0:2 = "torch.operator"(%arg0, %arg1, %arg2) <{name = "onnx.Loop"}> ({
^bb0(%arg4: !torch.vtensor<[],si64>, %arg5: !torch.vtensor<[],i1>, %arg6: !torch.vtensor<[],si32>):
  %1 = "torch.operator"(%arg6, %arg3) <{name = "onnx.Add"}> : (!torch.vtensor<[],si32>, !torch.vtensor<[],si32>) -> !torch.vtensor<[],si32>
  %2 = "torch.operator"(%arg6) <{name = "onnx.Identity"}> : (!torch.vtensor<[],si32>) -> !torch.vtensor<[],si32>
  %3 = "torch.operator"(%arg5) <{name = "onnx.Identity"}> : (!torch.vtensor<[],i1>) -> !torch.vtensor<[],i1>
  "torch.operator_terminator"(%3, %1, %2) : (!torch.vtensor<[],i1>, !torch.vtensor<[],si32>, !torch.vtensor<[],si32>) -> ()
}) : (!torch.vtensor<[],si64>, !torch.vtensor<[],i1>, !torch.vtensor<[],si32>) -> (!torch.vtensor<[],si32>, !torch.vtensor<[?],si32>)

Command

 iree-compile --iree-hal-target-backends=llvm-cpu --iree-llvmcpu-target-cpu=host -o test.vmfb model_torch_onnx.mlir 

Models Impacted:

ssd_mobilenet_v1_10
ssd_mobilenet_v1_12
ssd_mobilenet_v1_12-int8
ssd_mobilenet_v1_13-qdq
tiny-yolov3-11
yolov3-10
yolov3-12

[vivekkhandelwal1]

The issue with supporting this particular case is that when the Onnx.Loop op has scan outputs which are concatenated values after each iteration they can't be supported through a torch.prim.loop op. Since the scan output is a concatenated value of each iteration result the type of the Onnx.Loop op result and the value yielded through termination at each iteration is different resulting in errors in the downstream control-flow pipeline in LLVM.

Op to be lowered:

  %0:2 = torch.operator "onnx.Loop"(%arg0, %arg1, %arg2) : (!torch.vtensor<[],si64>, !torch.vtensor<[],i1>, !torch.vtensor<[],si32>) -> (!torch.vtensor<[],si32>, !torch.vtensor<[],si32>) {
    ^bb0(%arg4: !torch.vtensor<[],si64>, %arg5: !torch.vtensor<[],i1>, %arg6: !torch.vtensor<[],si32>):
      %592 = torch.operator "onnx.Add"(%arg6, %arg3) : (!torch.vtensor<[],si32>, !torch.vtensor<[],si32>) -> !torch.vtensor<[],si32>
      %593 = torch.operator "onnx.Identity"(%arg6) : (!torch.vtensor<[],si32>) -> !torch.vtensor<[],si32>
      %594 = torch.operator "onnx.Identity"(%arg5) : (!torch.vtensor<[],i1>) -> !torch.vtensor<[],i1>
      torch.operator_terminator %594, %592, %593 : !torch.vtensor<[],i1>, !torch.vtensor<[],si32>, !torch.vtensor<[],si32>
  }

Result IR (Generated after lowering Onnx.Loop):

    %8:2 = "torch.prim.Loop"(%0, %3, %arg2, %7) ({
    ^bb0(%arg7: !torch.int, %arg8: !torch.vtensor<[],si32>):
      %13 = "torch.prim.ListConstruct"() : () -> !torch.list<int>
      %14 = "torch.constant.none"() : () -> !torch.none
      %15 = "torch.constant.int"() <{value = 4 : i64}> : () -> !torch.int
      %16 = "torch.aten.full"(%13, %arg7, %15, %14, %14, %14) : (!torch.list<int>, !torch.int, !torch.int, !torch.none, !torch.none, !torch.none) -> !torch.vtensor<[],si64>
      %17 = "torch.operator"(%arg8, %arg3) <{name = "onnx.Add"}> : (!torch.vtensor<[],si32>, !torch.vtensor<[],si32>) -> !torch.vtensor<[],si32>
      %18 = "torch.operator"(%arg8) <{name = "onnx.Identity"}> : (!torch.vtensor<[],si32>) -> !torch.vtensor<[],si32>
      "torch.prim.Loop.condition"(%3, %17, %18) : (!torch.bool, !torch.vtensor<[],si32>, !torch.vtensor<[],si32>) -> ()
    }) : (!torch.int, !torch.bool, !torch.vtensor<[],si32>, !torch.vtensor<[],si32>) -> (!torch.vtensor<[],si32>, !torch.vtensor<[],si32>)

In the above IR the mismatch in the types can be seen.