Bump torch-mlir to 140cad5 and update TorchOnnxToTorch conversion pip…

…eline (iree-org#18867)

This commit bumps torch-mlir to
llvm/torch-mlir@140cad5.

This commit also replaces the `TorchOnnxToTorchPass` used to convert
torch-onnx IR to torch IR with `TorchOnnxToTorchBackendPipeline` which
was introduced here:
llvm/torch-mlir@fa4794d.
This pipeline consists of passes like `TorchOnnxToTorch`,
`ScalarizeShapes`, `DecomposeComplexOps`, `ShapeRefinementPipeline`,
etc., to convert the torch-onnx IR to torch IR consistent with the other
lowering paths for torch backend IR.

With the changes made in this PR, some of the tests which were failing
earlier during compilation now passes compilation, some of them even
passes inference, while there are some tests which were passing earlier
are now failing. The list is as follows:
- For
tests/external/iree-test-suites/onnx_ops/onnx_ops_cpu_llvm_sync.json:
  a.) Earlier failing compilation, now passing inference:
    - onnx/node/generated/test_convtranspose_kernel_shape
    - onnx/node/generated/test_einsum_sum

b.) Earlier failing compilation, now passing compilation but failing
inference
    - onnx/node/generated/test_convtranspose_output_shape
   
  c.) Earlier passing but now failing inference
    - onnx/node/generated/test_scan_sum
    - onnx/node/generated/test_scan9_sum

- For
tests/external/iree-test-suites/onnx_ops/onnx_ops_gpu_rocm_rdna3.json:
  a.) Earlier failing compilation, now passing inference:
    - onnx/node/generated/test_convtranspose_kernel_shape
    - onnx/node/generated/test_einsum_sum

b.) Earlier failing compilation, now passing compilation but failing
inference
    - onnx/node/generated/test_convtranspose_output_shape
   
  c.) Earlier passing but now failing inference
    - onnx/node/generated/test_scan_sum
    - onnx/node/generated/test_scan9_sum
    - onnx/node/generated/test_slice_default_axes
 
- For tests/external/iree-test-suites/onnx_ops/onnx_ops_gpu_vulkan.json:
  a.) Earlier failing compilation, now passing inference:
    - onnx/node/generated/test_convtranspose_kernel_shape

b.) Earlier failing compilation, now passing compilation but failing
inference
    - onnx/node/generated/test_convtranspose_output_shape
    - onnx/node/generated/test_einsum_sum

  c.) Earlier passing but now failing inference
    - onnx/node/generated/test_scan_sum
    - onnx/node/generated/test_scan9_sum
 
This commit
llvm/torch-mlir@55ff110
is expected to fix the newly introduced failures which will be included
in IREE in the next Torch-MLIR bump.

---------

Signed-off-by: Vivek Khandelwal <[email protected]>

compiler/plugins/input/Torch/PluginRegistration.cpp

@@ -57,8 +57,19 @@ struct TorchSession
       OpPassManager &passManager, std::string_view typeMnemonic) override {
     if (typeMnemonic == "onnx") {
       // ONNX input is a pre-processing step to torch.
-      passManager.addNestedPass<func::FuncOp>(
-          mlir::torch::onnx_c::createTorchOnnxToTorchPass());
+      mlir::torch::Torch::TorchLoweringPipelineOptions torchOnnxPipelineOptions;
+      // The `aten.flatten.using_ints` and `aten.unflatten.int` are added to the
+      // list of backend legal ops so that they are not decomposed into the
+      // `aten.view` op during the run of `DecomposeComplexOps` pass. The issue
+      // with this is that the `aten.view` op eventually lowers to
+      // `tensor.reshape` op while there exists a direct torch->linalg lowering
+      // for both the flatten/unflatten ops which lowers to
+      // `tensor.collapse_shape/expand_shape` op, and this is a more preferred
+      // path for the downstream pipeline.
+      torchOnnxPipelineOptions.backendLegalOps = {"aten.flatten.using_ints",
+                                                  "aten.unflatten.int"};
+      mlir::torch::Torch::createTorchOnnxToTorchBackendPipeline(
+          passManager, torchOnnxPipelineOptions);
     }
 
     if (typeMnemonic == "torch" || typeMnemonic == "onnx") {

tests/external/iree-test-suites/onnx_ops/onnx_ops_cpu_llvm_sync.json

@@ -101,8 +101,6 @@
     "onnx/node/generated/test_compress_default_axis",
     "onnx/node/generated/test_compress_negative_axis",
     "onnx/node/generated/test_convtranspose_autopad_same",
-    "onnx/node/generated/test_convtranspose_kernel_shape",
-    "onnx/node/generated/test_convtranspose_output_shape",
     "onnx/node/generated/test_cumsum_1d",
     "onnx/node/generated/test_cumsum_1d_exclusive",
     "onnx/node/generated/test_cumsum_1d_reverse",
@@ -127,7 +125,6 @@
     "onnx/node/generated/test_dft_inverse_opset19",
     "onnx/node/generated/test_dft_opset19",
     "onnx/node/generated/test_edge_pad",
-    "onnx/node/generated/test_einsum_sum",
     "onnx/node/generated/test_gridsample_bicubic",
     "onnx/node/generated/test_gridsample_bicubic_align_corners_0_additional_1",
     "onnx/node/generated/test_gridsample_bicubic_align_corners_1_additional_1",
@@ -427,6 +424,7 @@
     "onnx/node/generated/test_constantofshape_float_ones",
     "onnx/node/generated/test_constantofshape_int_shape_zero",
     "onnx/node/generated/test_constantofshape_int_zeros",
+    "onnx/node/generated/test_convtranspose_output_shape",
     "onnx/node/generated/test_dropout_default_mask_ratio",
     "onnx/node/generated/test_gridsample_nearest",
     "onnx/node/generated/test_gridsample_nearest_align_corners_0_additional_1",
@@ -447,6 +445,8 @@
     "onnx/node/generated/test_reduce_min_empty_set",
     "onnx/node/generated/test_reduce_sum_empty_set_non_reduced_axis_zero",
     "onnx/node/generated/test_resize_downsample_scales_linear_align_corners",
+    "onnx/node/generated/test_scan_sum",
+    "onnx/node/generated/test_scan9_sum",
     "onnx/node/generated/test_shape_clip_start",
     "onnx/node/generated/test_shape_end_1",
     "onnx/node/generated/test_shape_start_1",

tests/external/iree-test-suites/onnx_ops/onnx_ops_gpu_rocm_rdna3.json

@@ -105,8 +105,6 @@
     "onnx/node/generated/test_compress_default_axis",
     "onnx/node/generated/test_compress_negative_axis",
     "onnx/node/generated/test_convtranspose_autopad_same",
-    "onnx/node/generated/test_convtranspose_kernel_shape",
-    "onnx/node/generated/test_convtranspose_output_shape",
     "onnx/node/generated/test_cumsum_1d",
     "onnx/node/generated/test_cumsum_1d_exclusive",
     "onnx/node/generated/test_cumsum_1d_reverse",
@@ -131,7 +129,6 @@
     "onnx/node/generated/test_dft_inverse_opset19",
     "onnx/node/generated/test_dft_opset19",
     "onnx/node/generated/test_edge_pad",
-    "onnx/node/generated/test_einsum_sum",
     "onnx/node/generated/test_gridsample_bicubic",
     "onnx/node/generated/test_gridsample_bicubic_align_corners_0_additional_1",
     "onnx/node/generated/test_gridsample_bicubic_align_corners_1_additional_1",
@@ -442,6 +439,7 @@
     "onnx/node/generated/test_constantofshape_float_ones",
     "onnx/node/generated/test_constantofshape_int_shape_zero",
     "onnx/node/generated/test_constantofshape_int_zeros",
+    "onnx/node/generated/test_convtranspose_output_shape",
     "onnx/node/generated/test_dropout_default_mask_ratio",
     "onnx/node/generated/test_eyelike_populate_off_main_diagonal",
     "onnx/node/generated/test_eyelike_with_dtype",
@@ -490,6 +488,8 @@
     "onnx/node/generated/test_reduce_sum_square_default_axes_keepdims_random",
     "onnx/node/generated/test_reduce_sum_square_default_axes_keepdims_random_expanded",
     "onnx/node/generated/test_resize_downsample_scales_linear_align_corners",
+    "onnx/node/generated/test_scan_sum",
+    "onnx/node/generated/test_scan9_sum",
     "onnx/node/generated/test_shape",
     "onnx/node/generated/test_shape_clip_end",
     "onnx/node/generated/test_shape_clip_start",
@@ -501,6 +501,7 @@
     "onnx/node/generated/test_shape_start_negative_1",
     "onnx/node/generated/test_size",
     "onnx/node/generated/test_size_example",
+    "onnx/node/generated/test_slice_default_axes",
     "onnx/node/generated/test_split_zero_size_splits_opset13",
     "onnx/node/generated/test_split_zero_size_splits_opset18",
     "onnx/node/generated/test_top_k",

tests/external/iree-test-suites/onnx_ops/onnx_ops_gpu_vulkan.json

@@ -136,8 +136,6 @@
     "onnx/node/generated/test_compress_default_axis",
     "onnx/node/generated/test_compress_negative_axis",
     "onnx/node/generated/test_convtranspose_autopad_same",
-    "onnx/node/generated/test_convtranspose_kernel_shape",
-    "onnx/node/generated/test_convtranspose_output_shape",
     "onnx/node/generated/test_cumsum_1d",
     "onnx/node/generated/test_cumsum_1d_exclusive",
     "onnx/node/generated/test_cumsum_1d_reverse",
@@ -162,7 +160,6 @@
     "onnx/node/generated/test_dft_inverse_opset19",
     "onnx/node/generated/test_dft_opset19",
     "onnx/node/generated/test_edge_pad",
-    "onnx/node/generated/test_einsum_sum",
     "onnx/node/generated/test_gridsample",
     "onnx/node/generated/test_gridsample_aligncorners_true",
     "onnx/node/generated/test_gridsample_bicubic",
@@ -528,6 +525,7 @@
     "onnx/node/generated/test_constantofshape_int_zeros",
     "onnx/node/generated/test_convinteger_with_padding",
     "onnx/node/generated/test_convinteger_without_padding",
+    "onnx/node/generated/test_convtranspose_output_shape",
     "onnx/node/generated/test_dequantizelinear_int16",
     "onnx/node/generated/test_dequantizelinear_uint16",
     "onnx/node/generated/test_dropout_default_mask_ratio",
@@ -536,6 +534,7 @@
     "onnx/node/generated/test_einsum_batch_diagonal",
     "onnx/node/generated/test_einsum_batch_matmul",
     "onnx/node/generated/test_einsum_transpose",
+    "onnx/node/generated/test_einsum_sum",
     "onnx/node/generated/test_eyelike_with_dtype",
     "onnx/node/generated/test_isinf_float16",
     "onnx/node/generated/test_isnan_float16",
@@ -596,6 +595,8 @@
     "onnx/node/generated/test_reduce_sum_square_default_axes_keepdims_example_expanded",
     "onnx/node/generated/test_reduce_sum_square_default_axes_keepdims_random",
     "onnx/node/generated/test_reduce_sum_square_default_axes_keepdims_random_expanded",
+    "onnx/node/generated/test_scan_sum",
+    "onnx/node/generated/test_scan9_sum",
     "onnx/node/generated/test_shape_clip_start",
     "onnx/node/generated/test_shape_end_1",
     "onnx/node/generated/test_shape_start_1",