Onnx-Loop

[saienduri]

We are running into this error with onnx loop when running these models (dlrm). Repro instructions can be found here.

dlrm.default.pytorch.torch.mlir:488:12: error: failed to legalize operation 'torch.operator' that was explicitly marked illegal
    %378 = torch.operator "onnx.Loop"(%377, %328) : (!torch.vtensor<[],si64>, !torch.vtensor<[],i1>) -> !torch.vtensor<[?,16],f32> {
           ^
dlrm.default.pytorch.torch.mlir:488:12: note: see current operation: 
%187 = "torch.operator"(%54, %64) <{name = "onnx.Loop"}> ({
^bb0(%arg3: !torch.vtensor<[],si64>, %arg4: !torch.vtensor<[],i1>):
  %375 = "torch.aten.lt.Scalar"(%arg3, %109) : (!torch.vtensor<[],si64>, !torch.int) -> !torch.vtensor<[],i1>
  %376 = "torch.aten.add.Scalar"(%arg3, %53, %110) : (!torch.vtensor<[],si64>, !torch.int, !torch.int) -> !torch.vtensor<[],si64>
  %377 = "torch.aten.where.self"(%375, %376, %arg3) : (!torch.vtensor<[],i1>, !torch.vtensor<[],si64>, !torch.vtensor<[],si64>) -> !torch.vtensor<[],si64>
  %378 = "torch.aten.unsqueeze"(%377, %109) : (!torch.vtensor<[],si64>, !torch.int) -> !torch.vtensor<[1],si64>
  %379 = "torch.aten.index_select"(%185, %109, %378) : (!torch.vtensor<[129],si64>, !torch.int, !torch.vtensor<[1],si64>) -> !torch.vtensor<[1],si64>
  %380 = "torch.aten.squeeze"(%379) : (!torch.vtensor<[1],si64>) -> !torch.vtensor<[],si64>
  %381 = "torch.aten.lt.Scalar"(%arg3, %109) : (!torch.vtensor<[],si64>, !torch.int) -> !torch.vtensor<[],i1>
  %382 = "torch.aten.add.Scalar"(%arg3, %62, %110) : (!torch.vtensor<[],si64>, !torch.int, !torch.int) -> !torch.vtensor<[],si64>
  %383 = "torch.aten.where.self"(%381, %382, %arg3) : (!torch.vtensor<[],i1>, !torch.vtensor<[],si64>, !torch.vtensor<[],si64>) -> !torch.vtensor<[],si64>
  %384 = "torch.aten.unsqueeze"(%383, %109) : (!torch.vtensor<[],si64>, !torch.int) -> !torch.vtensor<[1],si64>
  %385 = "torch.aten.index_select"(%186, %109, %384) : (!torch.vtensor<[128],si64>, !torch.int, !torch.vtensor<[1],si64>) -> !torch.vtensor<[1],si64>
  %386 = "torch.aten.squeeze"(%385) : (!torch.vtensor<[1],si64>) -> !torch.vtensor<[],si64>
  %387 = "torch.aten.unsqueeze"(%380, %109) : (!torch.vtensor<[],si64>, !torch.int) -> !torch.vtensor<[1],si64>
  %388 = "torch.aten.unsqueeze"(%386, %109) : (!torch.vtensor<[],si64>, !torch.int) -> !torch.vtensor<[1],si64>
  %389 = "torch.aten.item"(%387) : (!torch.vtensor<[1],si64>) -> !torch.int
  %390 = "torch.aten.item"(%388) : (!torch.vtensor<[1],si64>) -> !torch.int
  %391 = "torch.aten.slice.Tensor"(%131, %109, %389, %390, %110) : (!torch.vtensor<[128],si64>, !torch.int, !torch.int, !torch.int, !torch.int) -> !torch.vtensor<[?],si64>
  %392 = "torch.aten.lt.Scalar"(%391, %109) : (!torch.vtensor<[?],si64>, !torch.int) -> !torch.vtensor<[?],i1>
  %393 = "torch.aten.add.Scalar"(%391, %52, %110) : (!torch.vtensor<[?],si64>, !torch.int, !torch.int) -> !torch.vtensor<[?],si64>
  %394 = "torch.aten.where.self"(%392, %393, %391) : (!torch.vtensor<[?],i1>, !torch.vtensor<[?],si64>, !torch.vtensor<[?],si64>) -> !torch.vtensor<[?],si64>
  %395 = "torch.aten.size.int"(%394, %109) : (!torch.vtensor<[?],si64>, !torch.int) -> !torch.int
  %396 = "torch.aten.index_select"(%76, %109, %394) : (!torch.vtensor<[1460,16],f32>, !torch.int, !torch.vtensor<[?],si64>) -> !torch.vtensor<[?,16],f32>
  %397 = "torch.prim.ListConstruct"(%109) : (!torch.int) -> !torch.list<int>
  %398 = "torch.aten.sum.dim_IntList"(%396, %397, %63, %108) : (!torch.vtensor<[?,16],f32>, !torch.list<int>, !torch.bool, !torch.none) -> !torch.vtensor<[16],f32>
  "torch.operator_terminator"(%64, %398) : (!torch.vtensor<[],i1>, !torch.vtensor<[16],f32>) -> ()
}) : (!torch.vtensor<[],si64>, !torch.vtensor<[],i1>) -> !torch.vtensor<[?,16],f32>

[renxida]

Been working on this at #332

closing in favor of that one.