AIWareK: Compiling PyTorch Model for AI Processor Using MLIR Framework

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.