ONNX To LLVM Pass Backend using linalg #3349

kimm240 · 2025-12-10T07:31:35Z

Overview

The Linalg pipeline transforms ONNX models into executable LLVM IR through a series of dialect conversions and optimizations. This document provides a detailed breakdown of all passes applied in the pipeline, including their locations, levels, and transformations.

Pipeline Flow

ONNX IR → [Preprocessing] → Linalg → [Bufferization] → Loops → Affine/SCF → CF → LLVM IR

Pass Application Summary Table

Step	Pass Name	Level	Location	Role	Input → Output
Phase 1: ONNX → Linalg (`addONNXToLinalgPasses`)
1.1	`createConvertONNXToLinalg()`	Func	`src/Conversion/ONNXToLinalg/ConvertONNXToLinalg.cpp` `src/Compiler/CompilerPasses.cpp:243`	ONNX → Linalg conversion	`onnx.MatMul` → `linalg.matmul`
1.2	`ConvertONNXEntryPointToKrnlPass`	Module	`src/Compiler/CompilerPasses.cpp:251-278` (uses `src/Conversion/ONNXToKrnl/ConvertONNXToKrnl.cpp`)	Entry point conversion	`onnx.EntryPoint` → `krnl.EntryPoint`
1.3	`createOneShotBufferizePass()`	Module	MLIR standard library `src/Compiler/CompilerPasses.cpp:284`	Tensor → Memref	`tensor<...>` → `memref<...>`
1.4	`createCanonicalizerPass()`	Module	MLIR standard library `src/Compiler/CompilerPasses.cpp:287`	Canonicalization	Optimization
Phase 2: Linalg → Affine/SCF → CF (`addLinalgToAffinePasses`)
2.1	`createConvertLinalgToLoopsPass()`	Func	MLIR standard library `src/Compiler/CompilerPasses.cpp:297`	Linalg → Loops	`linalg.matmul` → `affine.for`
2.2	`createBufferLoopHoistingPass()`	Func	MLIR standard library `src/Compiler/CompilerPasses.cpp:304`	Memory allocation hoisting	Move allocations outside loops
2.3	`buildBufferDeallocationPipeline()`	Func	MLIR standard library `src/Compiler/CompilerPasses.cpp:306-307`	Memory deallocation	Insert `dealloc` operations
2.4	`createOptimizeAllocationLivenessPass()`	Func	MLIR standard library `src/Compiler/CompilerPasses.cpp:308`	Lifetime optimization	Optimize allocation lifetimes
2.5	`createConvertBufferizationToMemRefPass()`	Func	MLIR standard library `src/Compiler/CompilerPasses.cpp:309`	Bufferization → MemRef	Standardize to memref ops
2.6	`createLowerAffinePass()`	Func	MLIR standard library `src/Compiler/CompilerPasses.cpp:312`	Affine → SCF	`affine.for` → `scf.for`
2.7	`createSCFToControlFlowPass()`	Func	MLIR standard library `src/Compiler/CompilerPasses.cpp:313`	SCF → CF	`scf.for` → `cf.br`
Phase 3: CF → LLVM (`addLinalgToLLVMPasses`)
3.1	`createConvertKrnlToLLVMPass()`	Module	`src/Conversion/KrnlToLLVM/ConvertKrnlToLLVM.cpp` `src/Compiler/CompilerPasses.cpp:334`	Krnl → LLVM + Runtime	`krnl.EntryPoint` → LLVM + Runtime functions
3.2	`createReconcileUnrealizedCastsPass()`	Module	MLIR standard library `src/Compiler/CompilerPasses.cpp:339`	Cast resolution	Resolve type casts
3.3	`createCanonicalizerPass()`	Module	MLIR standard library `src/Compiler/CompilerPasses.cpp:340`	Final canonicalization	Final optimization

Test

In test/mlir/conversion/onnx_to_linalg,

Test MLIR Code

Make MatMul.mlir in test/mlir/conversion/onnx_to_linalg
which has this code.

// Test MatMul conversion from ONNX to Linalg to LLVM IR
// This file tests the full pipeline: ONNX -> Linalg -> Affine -> LLVM IR

module {
  func.func @main_graph(%arg0: tensor<2x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<2x4xf32> {
    %0 = "onnx.MatMul"(%arg0, %arg1) : (tensor<2x3xf32>, tensor<3x4xf32>) -> tensor<2x4xf32>
    return %0 : tensor<2x4xf32>
  }
  "onnx.EntryPoint"() {func = @main_graph} : () -> ()
}

ONNX → LLVM dialect

export PATH=$PATH:$(pwd)/build/Release/bin && onnx-mlir --use-linalg-path --EmitLLVMIR MatMul.mlir -o test_linalg

LLVM dialect → .so

onnx-mlir -O3 --EmitLib test_linalg.onnx.mlir

Compile driver for check

g++ --std=c++11 driver.cpp -o driver -I ~/onnx-mlir/include -L ~/onnx-mlir/build/Release/lib -lcruntime -ldl

Run Driver

export LD_LIBRARY_PATH=~/onnx-mlir/build/Release/lib:$LD_LIBRARY_PATH && ./driver