[onnx] Enable --iree-input-type=onnx (#15995)
Includes a one-line fix to torch-mlir (already landed) that fixes an
unused variable warning.
With this, the following now works (for a model.onnx from the test
suite):
```
python -m iree.compiler.tools.import_onnx ~/tmp/model.onnx | ./tools/iree-compile --iree-input-type=onnx --iree-hal-target-backends=vmvx - -o /dev/null
```
We've still got a way to go on the conversions before there is enough
coverage to be generally usable.
diff --git a/compiler/plugins/input/Torch/torch-iree/CMakeLists.txt b/compiler/plugins/input/Torch/torch-iree/CMakeLists.txt
index b077221..4e5a220 100644
--- a/compiler/plugins/input/Torch/torch-iree/CMakeLists.txt
+++ b/compiler/plugins/input/Torch/torch-iree/CMakeLists.txt
@@ -20,6 +20,7 @@
torch-mlir::TorchConversionDialectIR
torch-mlir::TorchDialectIR
torch-mlir::TorchDialectPasses
+ torch-mlir::TorchOnnxToTorchPasses
torch-mlir::ConversionPasses
torch-mlir-dialects::TMTensorDialectIR
PUBLIC
diff --git a/compiler/plugins/input/Torch/torch-iree/PluginRegistration.cpp b/compiler/plugins/input/Torch/torch-iree/PluginRegistration.cpp
index 9f004e9..518a021 100644
--- a/compiler/plugins/input/Torch/torch-iree/PluginRegistration.cpp
+++ b/compiler/plugins/input/Torch/torch-iree/PluginRegistration.cpp
@@ -11,6 +11,7 @@
#include "torch-iree/InputConversion/Passes.h"
#include "torch-mlir-dialects/Dialect/TMTensor/IR/TMTensorDialect.h"
#include "torch-mlir/Conversion/Passes.h"
+#include "torch-mlir/Conversion/TorchOnnxToTorch/Passes.h"
#include "torch-mlir/Dialect/Torch/IR/TorchDialect.h"
#include "torch-mlir/Dialect/TorchConversion/IR/TorchConversionDialect.h"
#include "torch-mlir/Dialect/TorchConversion/Transforms/Passes.h"
@@ -39,6 +40,7 @@
mlir::torch::registerTorchPasses();
mlir::torch::registerTorchConversionPasses();
mlir::torch::registerConversionPasses();
+ mlir::torch::onnx_c::registerTorchOnnxToTorchPasses();
TorchInput::registerTMTensorConversionPasses();
}
@@ -51,7 +53,13 @@
bool extendCustomInputConversionPassPipeline(
OpPassManager &passManager, std::string_view typeMnemonic) override {
- if (typeMnemonic == "torch") {
+ if (typeMnemonic == "onnx") {
+ // ONNX input is a pre-processing step to torch.
+ passManager.addNestedPass<func::FuncOp>(
+ mlir::torch::onnx_c::createTorchOnnxToTorchPass());
+ }
+
+ if (typeMnemonic == "torch" || typeMnemonic == "onnx") {
TorchInput::TorchToIREELoweringPipelineOptions torchOptions;
torchOptions.strictSymbolicShapes = options.strictSymbolicShapes;
TorchInput::createTorchToIREEPipeline(passManager, torchOptions);
@@ -59,6 +67,7 @@
TorchInput::createConvertTMTensorToLinalgExtPass());
return true;
}
+
// TODO: Retire the tm_tensor input pipeline once we are fully switched
// to the 'torch' pipeline, which handles everything from the 'torch'
// dialect down (vs just 'tm_tensor' which was converting a couple of
@@ -74,6 +83,7 @@
void populateCustomInputConversionTypes(StringSet<> &typeMnemonics) override {
typeMnemonics.insert("tm_tensor");
typeMnemonics.insert("torch");
+ typeMnemonics.insert("onnx");
}
void populateDetectedCustomInputConversionTypes(
diff --git a/compiler/plugins/input/Torch/torch-mlir/CMakeLists.txt b/compiler/plugins/input/Torch/torch-mlir/CMakeLists.txt
index 5a1876b..c9c9302 100644
--- a/compiler/plugins/input/Torch/torch-mlir/CMakeLists.txt
+++ b/compiler/plugins/input/Torch/torch-mlir/CMakeLists.txt
@@ -200,6 +200,41 @@
)
###############################################################################
+# TorchOnnxToTorch
+###############################################################################
+
+file(GLOB _TorchOnnxToTorchPasses_SRCS
+ "${TORCH_MLIR_ROOT_DIR}/lib/Conversion/TorchOnnxToTorch/*.cpp"
+)
+iree_cc_library(
+ NAME
+ TorchOnnxToTorchPasses
+ SRCS
+ "${TORCH_MLIR_ROOT_DIR}/lib/Conversion/TorchOnnxToTorch/Passes.cpp"
+ ${_TorchOnnxToTorchPasses_SRCS}
+ DEPS
+ ::defs
+ ::TorchConversionDialectIR
+ ::TorchDialectIR
+ ::TorchOnnxToTorchPassesGen
+ MLIRArithDialect
+ MLIRFuncDialect
+ MLIRIR
+ MLIRPass
+)
+
+iree_tablegen_library(
+ NAME
+ TorchOnnxToTorchPassesGen
+ TD_FILE
+ "${TORCH_MLIR_ROOT_DIR}/include/torch-mlir/Conversion/TorchOnnxToTorch/Passes.td"
+ OUTS
+ -gen-pass-decls Conversion/TorchOnnxToTorch/Passes.h.inc
+ -gen-pass-capi-header Conversion/TorchOnnxToTorch/Passes.capi.h.inc
+ -gen-pass-capi-impl Conversion/TorchOnnxToTorch/Passes.capi.cpp.inc
+)
+
+###############################################################################
# CAPI
###############################################################################
diff --git a/third_party/torch-mlir b/third_party/torch-mlir
index be3e74b..ca631a5 160000
--- a/third_party/torch-mlir
+++ b/third_party/torch-mlir
@@ -1 +1 @@
-Subproject commit be3e74b647375d1c72f075bebf77dc31e74477a2
+Subproject commit ca631a52560acd6d80075034120d9a06ae2cb885
