AIWareK: Compiling PyTorch Model for AI Processor Using MLIR Framework
Hyunjeong Kwon, Hyun Mi Kim, Chun‐Gi Lyuh, Jin-Kyu Kim, Jinho Han, Young-Su Kwon · 2022 IEEE 4th International Conference on Artificial Intelligence Circuits and Systems (AICAS) · 2022
Deep learning compiler becomes necessary with the active research on AI hardware. This work compiles PyTorch models into target hardware codes using MLIR framework. The compiler first constructs a graph from the PyTorch model using TorchScript tracing. To construct the graph, our domain-specific parser generates an abstract syntax tree using tokens generated from the lexer. Then, the graph IR (GIR) is built and lowered into the kernel IR (KIR) and the processor IR (PIR) in MLIR framework. PIR becomes the input of the backed compiler that generates target machine codes. Experimental result shows that AIWareK compiled ResNet18 in 7.67 seconds, yielding 1.16e-03 mean absolute error.