Siddharth Jain

I am a PhD student and PMRF Fellow at Computer Science and Engeneering Department, Indian Institute of Technology Hyderabad.

My research interests are in the intersection of Compilers and Machine Learning. Currently I am working on Program Representation and Compiler Optimizations powered by ML/RL techniques.

Recent Activities

Mar 21, 2024	Received Kesav Nori Research Excellence Award for best MTech thesis in the system domain. Thesis title: Experiments with RL4ReAl: Benchmarking and Compile-Time Improvements
Mar 3, 2024	Gave a talk on “ML-Compiler-Bridge: Interfacing ML and Compilers”, at LLVM Performance Workshop, Co-located with International Symposium on Code Generation and Optimization (CGO)
Oct 26, 2023	Grateful to receive the Prime Minister’s Research Fellowship !
Jul 15, 2023	Completed MTech in Computer Science and Engneering from IIT Hyderabad!

Publications

The Next 700 ML-Enabled Compiler Optimizations

S. VenkataKeerthy, Siddharth Jain, Umesh Kalvakuntla, and 6 more authors

In Proceedings of the 33rd ACM SIGPLAN International Conference on Compiler Construction, 2024

Abs HTML PDF

There is a growing interest in enhancing compiler optimizations with ML models, yet interactions between compilers and ML frameworks remain challenging. Some optimizations require tightly coupled models and compiler internals, raising issues with modularity, performance and framework independence. Practical deployment and transparency for the end-user are also important concerns. We propose ML-Compiler-Bridge to enable ML model development within a traditional Python framework while making end-to-end integration with an optimizing compiler possible and efficient. We evaluate it on both research and production use cases, for training and inference, over several optimization problems, multiple compilers and its versions, and gym infrastructures.
RL4ReAl: Reinforcement Learning for Register Allocation

S. VenkataKeerthy, Siddharth Jain, Anilava Kundu, and 3 more authors

In Proceedings of the 32nd ACM SIGPLAN International Conference on Compiler Construction, 2023

Abs HTML PDF

We aim to automate decades of research and experience in register allocation, leveraging machine learning. We tackle this problem by embedding a multi-agent reinforcement learning algorithm within LLVM, training it with the state of the art techniques. We formalize the constraints that precisely define the problem for a given instruction-set architecture, while ensuring that the generated code preserves semantic correctness. We also develop a gRPC based framework providing a modular and efficient compiler interface for training and inference. Our approach is architecture independent: we show experimental results targeting Intel x86 and ARM AArch64. Our results match or out-perform the heavily tuned, production-grade register allocators of LLVM.