About Me

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.

AMD, India

Member of Technical Staff – AI Group

Oct 2023 – Present

Nod Labs, US

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.

Qualcomm, India

Machine Learning Compiler Engineer

July 2021 – Oct 2021

Education

Indian Institute of Science

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.

Madhav Institute of Technology and Science

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.