Abstract: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.

High Performance Code Generation in MLIR: An Early Case Study with GEMM

MLIR-based code generation for GPU tensor cores

Towards a high-performance AI compiler with upstream MLIR

Composable and Modular Code Generation in MLIR: A Structured and Retargetable Approach to Tensor Compiler Construction

The MLIR Transform Dialect. Your compiler is more powerful than you think

mlirSynth: Automatic, Retargetable Program Raising in Multi-Level IR using Program Synthesis

A Method for Efficient Heterogeneous Parallel Compilation: A Cryptography Case Study

NGEMM: Optimizing GEMM for Deep Learning via Compiler-based Techniques

MPIrigen: MPI Code Generation through Domain-Specific Language Models

Experiences Building an MLIR-based SYCL Compiler

HIR: An MLIR-based Intermediate Representation for Hardware Accelerator Description

MCompiler: A Synergistic Compilation Framework

Fortran performance optimisation and auto-parallelisation by leveraging MLIR-based domain specific abstractions in Flang

Deep Data Flow Analysis

Building a Reusable and Extensible Automatic Compiler Infrastructure for Reconfigurable Devices

Assessing Code Generation with Intermediate Languages

ML-driven Hardware Cost Model for MLIR

Generating performance portable code using rewrite rules: from high-level functional expressions to high-performance OpenCL code

CMLCompiler: A Unified Compiler for Classical Machine Learning

NonGEMM Bench: Understanding the Performance Horizon of the Latest ML Workloads with NonGEMM Workloads

Performance-Aligned LLMs for Generating Fast Code