Optimization-Preserving Lowering Pipeline Analysis for Detecting Missed Optimizations in MLIR-based AI Compilers

ABSTRACT

MLIR is widely used in AI compilers to progressively lower high-level tensor or domain-specific programs into lower-level representations for hardware execution. During this lowering process, compiler optimizations should ideally preserve or expose optimization opportunities rather than accidentally hide them.

This project explores a testing approach for detecting **missed optimization bugs** in MLIR lowering pipelines. The core idea is to construct or identify lowering pipelines that should be **optimization-preserving**: if an optimization can be applied before lowering, then an equivalent optimization opportunity should still be visible after lowering. The student will build a prototype framework that applies optimizations to programs before and after lowering, then compares whether both versions trigger the expected optimizations.

For example, given an MLIR program, the framework may run an optimization pass at a high-level dialect, lower the program to a lower-level dialect, and then run corresponding optimizations again. If the optimization is applied before lowering but not after lowering, this may indicate that the lowering pipeline obscures useful structure and causes a missed optimization.

The goal is to automatically detect cases where MLIR lowering introduces representations that are semantically correct but optimization-unfriendly, leading to inefficient generated code.

Research questions include:

• Can we define lowering pipelines that should preserve specific optimization opportunities?
• Can differential analysis before and after lowering reveal missed optimization bugs?
• Which MLIR dialects, operations, or lowering passes are most likely to hide optimization opportunities?
• Can this method find real missed optimizations in MLIR-based AI compiler pipelines?

SKILLS & EXPERIENCE REQUIRED

• Background in Computer Science or related discipline.
• Good programming skills, preferably in C++ and Python.
• Interest in LLMs, agents, or AI-assisted programming.
• Interest in compilers, AI infra, or software testing.
• Basic familiarity with LLVM/MLIR is helpful but not required.