Converting ONNX Models to StableHLO

[brataTT]

What is the proper way to convert onnx models to stablehlo?

I'm encountering issues converting ONNX models from the model zoo to StableHLO using onnx-mlir. While some models convert partially, others fail completely.

I have tried the following commands:

onnx-mlir --EmitONNXIR <model>.onnx
onnx-mlir-opt  --shape-inference --convert-onnx-to-stablehlo <model>.onnx.mlir > shlo.mlir

But I get partial stablehlo conversion. onnx.conv, onnx.QuantizeLinear and onnx.DequantizeLinear remain in the stablehlo output.
After checking the list of supported ops and the opsets used by the model zoo for these models. It seems the operations in question should be supported.

Example outputs I get for resnet models

resnet18-v1-7.onnx:

wget https://github.com/onnx/models/raw/refs/heads/main/validated/vision/classification/resnet/model/resnet18-v1-7.onnx
onnx-mlir --EmitONNXIR resnet18-v1-7.onnx
onnx-mlir-opt  --shape-inference --convert-onnx-to-stablehlo resnet18-v1-7.onnx.mlir > shlo.mlir

onnx.conv is still present in the stablehlo output

%0 = "onnx.Conv"(%arg0, %cst_39, %cst_40) {auto_pad = "NOTSET", dilations = [1, 1], group = 1 : si64, kernel_shape = [7, 7], onnx_node_name = "resnetv15_batchnorm0_fwd-resnetv15_conv0_fwd_0", pads = [3, 3, 3, 3], strides = [2, 2]} : (tensor<?x3x224x224xf32>, tensor<64x3x7x7xf32>, tensor<64xf32>) -> tensor<?x64x112x112xf32>

resnet50-v1-12-qdq.onnx:

wget https://github.com/onnx/models/raw/refs/heads/main/validated/vision/classification/resnet/model/resnet50-v1-12-qdq.onnx
onnx-mlir --EmitONNXIR resnet50-v1-12-qdq.onnx
onnx-mlir-opt  --shape-inference --convert-onnx-to-stablehlo resnet50-v1-12-qdq.onnx.mlir > shlo.mlir

onnx.QuantizeLinear and onnx.DequantizeLinear are still present in the stablehlo output

%0 = "onnx.QuantizeLinear"(%arg0, %cst_18, %c_17) {axis = 1 : si64, onnx_node_name = "data_resnetv17_conv0_fwd_QuantizeLinear", saturate = 1 : si64} : (tensor<?x3x224x224xf32>, tensor<f32>, tensor<ui8>) -> tensor<?x3x224x224xui8>
%1 = "onnx.DequantizeLinear"(%c_24, %cst_25, %c_26) {axis = 1 : si64, onnx_node_name = "ConvBnFusion_BN_B_resnetv17_batchnorm0_beta_DequantizeLinear"} : (tensor<64xi32>, tensor<1xf32>, tensor<1xi32>) -> tensor<64xf32>

resnet50-v1-12-int8.onnx:

wget https://github.com/onnx/models/raw/refs/heads/main/validated/vision/classification/resnet/model/resnet50-v1-12-int8.onnx
onnx-mlir --EmitONNXIR resnet50-v1-12-int8.onnx 
onnx-mlir-opt  --shape-inference --convert-onnx-to-stablehlo resnet50-v1-12-int8.onnx.mlir > shlo.mlir 

The model fails to convert to stablehlo with the following error:

#0 0x00000040040020e7 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (/workdir/onnx-mlir/build/Debug/bin/onnx-mlir-opt+0x40020e7)
 #1 0x0000004003fffc7e llvm::sys::RunSignalHandlers() (/workdir/onnx-mlir/build/Debug/bin/onnx-mlir-opt+0x3fffc7e)
 #2 0x000000400400279a SignalHandler(int) Signals.cpp:0:0
 #3 0x0000004008950520 (/lib/x86_64-linux-gnu/libc.so.6+0x42520)
 #4 0x00000040089a49fc pthread_kill (/lib/x86_64-linux-gnu/libc.so.6+0x969fc)
 #5 0x0000004008950476 gsignal (/lib/x86_64-linux-gnu/libc.so.6+0x42476)
 #6 0x00000040089367f3 abort (/lib/x86_64-linux-gnu/libc.so.6+0x287f3)
 #7 0x000000400893671b (/lib/x86_64-linux-gnu/libc.so.6+0x2871b)
 #8 0x0000004008947e96 (/lib/x86_64-linux-gnu/libc.so.6+0x39e96)
 #9 0x0000004000ac0958 decltype(auto) llvm::cast<mlir::FloatType, mlir::Type>(mlir::Type&) /workdir/llvm-project/llvm/include/llvm/Support/Casting.h:573:37
#10 0x0000004001222f11 onnx_mlir::(anonymous namespace)::createInitialValueForPoolingOp(mlir::Operation*, mlir::Type, mlir::ConversionPatternRewriter&) /workdir/onnx-mlir/src/Conversion/ONNXToStablehlo/NN/Pooling.cpp:32:29
#11 0x0000004001221c80 onnx_mlir::(anonymous namespace)::ONNXPoolOpLoweringToStablehlo<mlir::ONNXMaxPoolSingleOutOp, mlir::ONNXMaxPoolSingleOutOpAdaptor, onnx_mlir::ONNXGenericPoolOpShapeHelper<mlir::ONNXMaxPoolSingleOutOp>>::matchAndRewrite(mlir::Operation*, llvm::ArrayRef<mlir::Value>, mlir::ConversionPatternRewriter&) const /workdir/onnx-mlir/src/Conversion/ONNXToStablehlo/NN/Pooling.cpp:131:21
#12 0x0000004000413b56 mlir::ConversionPattern::matchAndRewrite(mlir::Operation*, llvm::ArrayRef<mlir::ValueRange>, mlir::ConversionPatternRewriter&) const /workdir/llvm-project/mlir/include/mlir/Transforms/DialectConversion.h:564:12
#13 0x0000004003036241 mlir::ConversionPattern::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&) const (/workdir/onnx-mlir/build/Debug/bin/onnx-mlir-opt+0x3036241)
#14 0x00000040030784a9 void llvm::function_ref<void ()>::callback_fn<mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<llvm::LogicalResult (mlir::Pattern const&)>)::$_2>(long) PatternApplicator.cpp:0:0
#15 0x0000004003074daf mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<llvm::LogicalResult (mlir::Pattern const&)>) (/workdir/onnx-mlir/build/Debug/bin/onnx-mlir-opt+0x3074daf)
#16 0x00000040030372f2 (anonymous namespace)::OperationLegalizer::legalize(mlir::Operation*, mlir::ConversionPatternRewriter&) DialectConversion.cpp:0:0
#17 0x0000004003036367 mlir::OperationConverter::convert(mlir::ConversionPatternRewriter&, mlir::Operation*) (/workdir/onnx-mlir/build/Debug/bin/onnx-mlir-opt+0x3036367)
#18 0x000000400303750f mlir::OperationConverter::convertOperations(llvm::ArrayRef<mlir::Operation*>) (/workdir/onnx-mlir/build/Debug/bin/onnx-mlir-opt+0x303750f)
#19 0x000000400303dc6b mlir::applyPartialConversion(mlir::Operation*, mlir::ConversionTarget const&, mlir::FrozenRewritePatternSet const&, mlir::ConversionConfig) (/workdir/onnx-mlir/build/Debug/bin/onnx-mlir-opt+0x303dc6b)
#20 0x00000040011bb624 onnx_mlir::FrontendToStablehloLoweringPass::runOnOperation() /workdir/onnx-mlir/src/Conversion/ONNXToStablehlo/ConvertONNXToStablehlo.cpp:133:14
#21 0x0000004003cd77c4 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (/workdir/onnx-mlir/build/Debug/bin/onnx-mlir-opt+0x3cd77c4)
#22 0x0000004003cd7f01 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) (/workdir/onnx-mlir/build/Debug/bin/onnx-mlir-opt+0x3cd7f01)
#23 0x0000004003cda5ab mlir::PassManager::run(mlir::Operation*) (/workdir/onnx-mlir/build/Debug/bin/onnx-mlir-opt+0x3cda5ab)
#24 0x000000400143425f performActions(llvm::raw_ostream&, std::shared_ptr<llvm::SourceMgr> const&, mlir::MLIRContext*, mlir::MlirOptMainConfig const&) MlirOptMain.cpp:0:0
#25 0x0000004001433d88 llvm::LogicalResult llvm::function_ref<llvm::LogicalResult (std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>::callback_fn<mlir::MlirOptMain(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, mlir::DialectRegistry&, mlir::MlirOptMainConfig const&)::$_3>(long, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&) MlirOptMain.cpp:0:0
#26 0x0000004003f26705 mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<llvm::LogicalResult (std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>, llvm::raw_ostream&, llvm::StringRef, llvm::StringRef)::$_0::operator()(llvm::StringRef) const ToolUtilities.cpp:0:0
#27 0x0000004003f26372 mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<llvm::LogicalResult (std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>, llvm::raw_ostream&, llvm::StringRef, llvm::StringRef) (/workdir/onnx-mlir/build/Debug/bin/onnx-mlir-opt+0x3f26372)
#28 0x000000400142d061 mlir::MlirOptMain(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, mlir::DialectRegistry&, mlir::MlirOptMainConfig const&) (/workdir/onnx-mlir/build/Debug/bin/onnx-mlir-opt+0x142d061)
#29 0x00000040002caf96 main /workdir/onnx-mlir/src/Tools/onnx-mlir-opt/onnx-mlir-opt.cpp:209:14
#30 0x0000004008937d90 (/lib/x86_64-linux-gnu/libc.so.6+0x29d90)
#31 0x0000004008937e40 __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x29e40)
#32 0x00000040002c9c55 _start (/workdir/onnx-mlir/build/Debug/bin/onnx-mlir-opt+0x2c9c55)
Aborted (core dumped)

[AlexandreEichenberger]

@brataTT Sorry we don't use this path much at IBM. However, we would welcome to add a target like -EmitStableHLOIR that would be a standard driver. If you get the proper sequence of onnx-mlir-opt, we can do that work.

[brataTT]

Thank you for the response @AlexandreEichenberger
Are there any docs I can check to help understand what is currently possible/supported?
I'm checking SupportedONNXOps-cpu.md to figure out which ops I can and cannot expect to be supported, but that didn't help me.

[AlexandreEichenberger]

My best suggestion would be to look at recent StableHLO PRs and ping the authors, as I have no visibility on how they use onnx-mlir with drivers outside of this repo.

I'm checking SupportedONNXOps-cpu.md to figure out which ops I can and cannot expect to be supported, but that didn't help me.

That file has the ops that we support when going all the way down to binary, StableHLO is a different path and the support provided essentially by the conversion of ONNX ops to StableHLO does not correlate with what is lowered to binary.