  During early optimization passes, compilers must make predictions for
machine-dependent characteristics such as execution unit utilization, number of
register spills, latency, throughput etc. to generate better code. Often a
hand-written static/analytical hardware cost model is built into the compiler.
However, the need for more sophisticated and varied predictions has become more
pronounced with the development of deep learning compilers which need to
optimize dataflow graphs. Such compilers usually employ a much higher level
MLIR form as an IR representation before lowering to traditional LLVM-IR. A
static/analytical cost model in such a scenario is cumbersome and error prone
as the opcodes represent very high level algebraic/arithmetic operations.
Hence, we develop a machine learning-based cost model for high-level MLIR which
can predict different target variables of interest such as CPU/GPU/xPU
utilization, instructions executed, register usage etc. By considering the
incoming MLIR as a text input a la NLP models we can apply well-known
techniques from modern NLP research to help predict hardware characteristics
more accurately. We expect such precise ML-driven hardware cost models to guide
our deep learning compiler in graph level optimizations around operator fusion,
local memory allocation, kernel scheduling etc. as well as in many kernel-level
optimizations such as loop interchange, LICM and unroll. We report early
work-in -progress results of developing such models on high-level MLIR
representing dataflow graphs emitted by Pytorch/Tensorflow-like frameworks as
well as lower-level dialects like affine. 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.
