Control-flow Aware MLIR Tracing.

MLIR (Multi-Level Intermediate Representation) has proven to be a scalable framework for building modern compilers. A key challenge, however, lies in effectively provide these compilers with input programs. Several methods have been explored: static methods, such as Domain-Specific Language, generic language front-end that transform an abstract syntax tree using tools like Polygeist, or leveraging compiler intermediate representation with Brutus. Dynamic approaches, such as tracing with JAX or Reactant, offer another solution. Each class of methods presents trade-offs: static approaches are easily integrated through compiler passes but lack runtime information, while dynamic tracing captures runtime behavior but cannot handle compile-time information. In this paper, we propose enhancing the Reactant tracing system by integrating the Julia compiler static analysis. This hybrid approach improves the expressiveness and precision of tracing. The effectiveness of our approach is demonstrated on the Polybench benchmark suite.