I am an Artificial Intelligence Framework Engineer and work at the intersection of Artificial Intelligence and High-Performance Computing. I primarily work on developing the frontend for AI/ML models for computer vision and natural language processing for seamless integration on different sets of hardware. Most of my work is based on the PyTorch and MLIR compiler infrastructure. I have also worked in the area of automatic code generation for GPUs. I actively contribute to various open-source projects like Torch-MLIR, MLIR, IREE, PyTorch, etc.
Work
Here, I work on developing frontend support for ONNX AI models to ensure seamless integration with AMD processors, enabling efficient execution of advanced AI workloads. Implemented model optimization techniques to enhance performance and reduce computational overhead.
Member of Technical Staff – AI Group
Oct 2023 – Present
AI Compiler Engineer
Nov 2021 – Oct 2023
Here, my work was focused on enabling and optimizing state-of-the-art AI/ML models, including BERT, ResNet, GPT-2, and Stable Diffusion, for both training and
inference across diverse hardware platforms such as NVIDIA GPUs, AMD GPUs, and edge devices using the MLIR compiler infrastructure and IREE runtime.
Here, I worked on the Glow compiler infrastructure. The task was to enable support for the Deep Learning models like MobileNet, FaceNet, etc., on the company’s hardware and enhance their performance.
Machine Learning Compiler Engineer
July 2021 – Oct 2021
Education
MTech in CSE
July 2019 – June 2021
Here, I joined the Multicore Computing Lab and did my thesis in the area of Automatic code generation for GPU Tensor Cores using MLIR under the guidance of Prof. Uday Reddy Bondhugula. I also studied courses like Machine Learning, Deep Learning for NLP, Cryptography, Computer Security, Compiler, and Practical Data Science.
I studied the foundational courses of computer science engineering from streams like theory and software systems and did my major project under the guidance of Prof. Jaimala Jha. I also organized various technical events as the vice president of the institute’s Digital Learning Club.
BE in CSE
July 2015 – June 2019
Publication
I have a research paper titled “MLIR-Based Code Generation for GPU Tensor Cores” published at the 31st ACM SIGPLAN INTERNATIONAL CONFERENCE ON COMPILER CONSTRUCTION (CC’22). The work presents an MLIR-based code-generation pipeline for Tensor Cores on NVIDIA GPUS, which achieves performance comparable to that of hand-tuned libraries like cuBLAS, cuBLASLt, and cuTensor.