Reproduce “Compiling ONNX Neural Network Models Using MLIR” on OpenResearch

About this paper

Tian Jin, Gheorghe-Teodor Bercea, Tung D. Le, Tong Chen, Gong Su, Haruki Imai, Yasushi Negishi, Anh Leu and 3 more

Researchers from IBM developed `onnx-mlir`, an open-source compiler that translates ONNX deep neural network models into optimized native code using the MLIR infrastructure, enhancing model portability and inference efficiency across diverse hardware platforms like x86, IBM Power Systems, and IBM System Z. This compiler, covering 51 ONNX operations, supports critical models like MNIST and ResNet50.

Key insights

• MLIR serves as a powerful and extensible infrastructure for developing sophisticated deep learning compilers, effectively managing multi-level abstractions and declarative optimization rules.
• The custom `krnl` dialect effectively separates computation semantics from execution scheduling, establishing a robust intermediate representation for applying advanced loop-based and polyhedral optimizations to neural network operations.
• Mature compiler optimization techniques, including graph rewriting, constant propagation, and shape inference, can be systematically applied within the MLIR framework to significantly improve deep learning model inference efficiency.

Summary by alphaXiv.