Bump torch-mlir to 140cad5 and update TorchOnnxToTorch conversion pipeline (#18867) This commit bumps torch-mlir to https://github.com/llvm/torch-mlir/commit/140cad5659bb779bb1f5de1888566db5b5d21236. This commit also replaces the `TorchOnnxToTorchPass` used to convert torch-onnx IR to torch IR with `TorchOnnxToTorchBackendPipeline` which was introduced here: https://github.com/llvm/torch-mlir/commit/fa4794dae2057876ec8ad2a6464e2668f6a2ea0c. 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 https://github.com/llvm/torch-mlir/commit/55ff110dc29cab7e2495ccdbec9a60512c29c665 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 <vivekkhandelwal1424@gmail.com>
diff --git a/compiler/plugins/input/Torch/PluginRegistration.cpp b/compiler/plugins/input/Torch/PluginRegistration.cpp index b9a497f..6f79686 100644 --- a/compiler/plugins/input/Torch/PluginRegistration.cpp +++ b/compiler/plugins/input/Torch/PluginRegistration.cpp
@@ -57,8 +57,19 @@ 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") {
diff --git a/tests/external/iree-test-suites/onnx_ops/onnx_ops_cpu_llvm_sync.json b/tests/external/iree-test-suites/onnx_ops/onnx_ops_cpu_llvm_sync.json index 0bbd604..a025431 100644 --- a/tests/external/iree-test-suites/onnx_ops/onnx_ops_cpu_llvm_sync.json +++ b/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",
diff --git a/tests/external/iree-test-suites/onnx_ops/onnx_ops_gpu_rocm_rdna3.json b/tests/external/iree-test-suites/onnx_ops/onnx_ops_gpu_rocm_rdna3.json index 9901e01..79cc5a9 100644 --- a/tests/external/iree-test-suites/onnx_ops/onnx_ops_gpu_rocm_rdna3.json +++ b/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",
diff --git a/tests/external/iree-test-suites/onnx_ops/onnx_ops_gpu_vulkan.json b/tests/external/iree-test-suites/onnx_ops/onnx_ops_gpu_vulkan.json index e6d9a4a..8c31c26 100644 --- a/tests/external/iree-test-suites/onnx_ops/onnx_ops_gpu_vulkan.json +++ b/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",
diff --git a/third_party/torch-mlir b/third_party/torch-mlir index 45bb17e..140cad5 160000 --- a/third_party/torch-mlir +++ b/third_party/torch-mlir
@@ -1 +1 @@ -Subproject commit 45bb17ebfe5e9cdcfd2cfabf850d9dec7127c5ab +Subproject commit 140cad5659bb779bb1f5de1888566db5b5d21236