ML Driven Hardware Cost Model for MLIR | Online Only

ML Driven Hardware Cost Model for MLIR by Dr. Dibyendu Das

Abstract:
------------
Compiler optimizations are diverse and have different payoffs for different types of code generated. During high-level optimization passes, compilers must make predictions for machine-dependent characteristics such as execution unit utilization, number of register spills, latency, throughput etc. to compare various heuristics for improved code generation. However, the need for more sophisticated and varied predictions has become more pronounced with the development of Deep Learning (DL) Compilers (Tensorflow XLA, Pytorch 2.0 Dynamo/Inductor, TVM, Habana Synapse AI) which need to optimize dataflow graphs. Such compilers usually employ a much higher level MLIR as an IR representation before lowering to traditional LLVM-IR. A static/analytical cost model in such a scenario is cumbersome and error prone. We develop a machine learning (ML)-based cost model for MLIR which can predict different hardware target variables of interest for CPU/GPU/xPU. Such an approach is of great benefit to an AI SW stack which spans across x86 CPU, GPU, FPGA. In addition to the DL stack, these ML-driven HW cost models can also be used in traditional compilers which are slowly but surely moving to embrace MLIR. We have reported early work-in-progress results of building such models on high-level MLIR in LLVM-DEVELOPERS’ MEETING 2022 (https://llvm.org/devmtg/2022-11/slides/TechTalk13-ML-basedHardwareCostModel-MLIR.pdf). We show that these models can provide reasonably good estimates with low error bounds for various hardware characteristics of interest and can be a go-to mechanism for hardware cost modelling in the future. An arXiv version of the work can also be found in: https://arxiv.org/abs/2302.11405 .

Speaker's Bio:
--------------------
Dr. Dibyendu Das is a senior technologist working at the confluence of AI and compilers/tools. He works on Deep Learning compiler frameworks for efficient code generation as well as explore future technologies related to enablement of AI in compiler algorithms. This involves working with MLIR and LLVM for ML frameworks and underlying compilation technologies, DL hardware/accelerators and novel mechanisms for code generation/optimization. He was previously the lead of the team at AMD that developed AOCC - AMD's premium optimizing compiler that achieved world record SPEC benchmark scores.