onnx-mlir icon indicating copy to clipboard operation
onnx-mlir copied to clipboard

How to generate pure MLIR?

Open yuucyf opened this issue 1 year ago • 6 comments

When using onnx-mlir, the --EmitMLIR parameter is described as "Lower the input to MLIR built-in transformation dialect," which might lead one to believe it generates pure MLIR code. However, after actual use, it becomes clear that the generated file still contains krnl operations, such as krnl.global, krnl.entry_point, and so on.

So, does onnx-mlir provide a way to generate pure MLIR? Currently, I’m unsure how many krnl operations are included in the MLIR file generated by --EmitMLIR, and how these operations can be converted to MLIR built-in IR. Can anyone help clarify this?

yuucyf avatar Feb 20 '25 06:02 yuucyf

You are correct, we are using KRNL operations inside of the EmitMLIR target as some KRNL operations are lowered all the way to LLVM lowering. -EmitMLIR means we removed all of the ONNX dialect operation.

If you want to intersect with dialects only found in MLIR official repo, then you will need to lower to TOSA or StableHLO. This is a path that we at IBM do not use, so you will have to touch base with developers that use this path to get more info on this.

AlexandreEichenberger avatar Feb 20 '25 15:02 AlexandreEichenberger

You are correct, we are using KRNL operations inside of the EmitMLIR target as some KRNL operations are lowered all the way to LLVM lowering. -EmitMLIR means we removed all of the ONNX dialect operation.

If you want to intersect with dialects only found in MLIR official repo, then you will need to lower to TOSA or StableHLO. This is a path that we at IBM do not use, so you will have to touch base with developers that use this path to get more info on this.

KRNL operations are used in the EmitMLIR target. Is it possible to lower all KRNL operations in EmitMLIR target(or EmitPureMLIR target) to pure MLIR?

For example: krnl.global() => memref.global() & memref.get_global() ...

yuucyf avatar Mar 03 '25 03:03 yuucyf

I know for a fact that in our path, where we lower down all the way to LLVM dialect, there are certain Krnl Ops that can only be lowered at the same time as LLVM.

What you are suggesting is nevertheless interesting, namely that we we don't go that path and are targeting EmitMLIR we could lower some of these remaining Krnl ops earlier. I can check one or two to see if that is possible. Would you be willing @yuucyf to help a bit with this effort, as this is not something that we would be directly using?

AlexandreEichenberger avatar Mar 03 '25 17:03 AlexandreEichenberger

I know for a fact that in our path, where we lower down all the way to LLVM dialect, there are certain Krnl Ops that can only be lowered at the same time as LLVM.

What you are suggesting is nevertheless interesting, namely that we we don't go that path and are targeting EmitMLIR we could lower some of these remaining Krnl ops earlier. I can check one or two to see if that is possible. Would you be willing @yuucyf to help a bit with this effort, as this is not something that we would be directly using?

My current goal is to use onnx-mlir to convert ONNX model files into pure MLIR. I plan to modify the tool's workflow to achieve this. As I'm just starting with onnx-mlir, I might need some time to familiarize myself with the codebase. Would you recommend any specific entry points in the compilation pipeline where Krnl lowering could be decoupled from LLVM generation? I'd be happy to collaborate on prototyping this.

yuucyf avatar May 08 '25 06:05 yuucyf

The way I often look at things is to do a compilation of a small/medium model, and trace with onnx-mlir -mlir-print-ir-after-all. You can then see which pass does what. You can grep the pass name from the listing, and you then can locate the pass that does a transform. You can then change the code as you see fit and having it reviewed.

If you have a model that has large constants, you can elide them by adding the following additional options --mlir-elide-elementsattrs-if-larger=20 --mlir-elide-resource-strings-if-larger=20.

If you cat the output to a file with an .mlir ending, it let you look at the log in VSCode with the auto formatting. The script below cut the max line to a given length and a few things that make it looks better

onnx-mlir ${@:1:$#-1}  -mlir-print-ir-after-all --mlir-elide-elementsattrs-if-larger=20 --mlir-elide-resource-strings-if-larger=20 2>&1 | cut -c -800 | sed -E "s/^(([^\"]*\"[^\"]*\")*[^\"]*\"[^\"]*)$/\1  \"/" | tee ${@: -1}

If that script is named onnx-mlir-print.sh, the you can use it as follows `onnx-mlir-print.sh <options...> <log-file-name.mlir>, where the log file should be perfect for formatting in the VSCode with mlir formatting.

AlexandreEichenberger avatar May 08 '25 13:05 AlexandreEichenberger

I am developing a conversion pass from the ONNX dialect to MLIR's built-in dialects, such as Linalg, Tensor, and Arith. Support for a subset of operators has already been implemented.

jiangnan7 avatar Jul 29 '25 08:07 jiangnan7