Senior Software Engineer
Google
My Technical Blog
Google Scholar
Curriculum Vitae

I am a Senior Software Enngineer at Google working on the XLA compiler for TPUs with a focus on large-scale distributed training of ML workloads. I worked on the distributed execution infrastructure, scalable collective algorithms, scheduling of HLOs to overlap the communications with computations, and pipelining and fusion of HLOs. I worked with HLOs, MLIR, and LLVM-based Intermediate Representations. I also worked on optimizing the performance of various ML workloads on TPUs with a fous in embedding-based workloads. Other workloads include Neural Radiance Fields (NeRF) and Graph Neural Networks (GNN),

I was a PhD student in the COMMIT group at MIT CSAIL and advised by Prof. Saman Amarasinghe and Prof. Julian Shun. I completed my Bachelor and Master programs at Rice University advised by Prof. Vivek Sarkar and Prof. Alan L. Cox. I have developed the GraphIt domain-specific language and compiler for writing high-performance graph analytics applications. GraphIt is used by Nvidia, UW, and Cornell for the development of domain-specific accelerators. I have also worked on cache optimizations for graph analytics and sparse linear algebra

Publications

Taming the Zoo: The Unified GraphIt Compiler Framework for Novel Architectures
[pdf]
Ajay Brahmakshatriya, Emily Furst, Victor A. Ying, Claire Hsu, Changwan Hong, Max Ruttenberg, Yunming Zhang, Tommy Jung, Dustin Richmond, Michael Taylor, Julian Shun, Mark Oskin, Daniel Sanchez, and Saman Amarasinghe
Compiling Graph Algorithms for GPUs with GraphIt
[pdf] [arxiv]
Ajay Brahmakshatriya, Yunming Zhang, Changwan Hong, Shoaib Kamil, Saman Amarasinghe, and Julian Shun
International Symposium on Code Generation and Optimization (CGO 2021)
Best Paper Award
GraphIt: Optimizing the Performance and Improving the Programmability of Graph Algorithms
[pdf] [GraphIt project]
Yunming Zhang
PhD Thesis (Sep 2020)
Optimizing Ordered Graph Algorithms with GraphIt
[pdf] [arxiv] [GraphIt project]
Yunming Zhang, Ajay Brahmakshatriya, Xinyi Chen, Laxman Dhulipala, Shoaib Kamil, Saman Amarasinghe, and Julian Shun
International Symposium on Code Generation and Optimization (CGO 2020)
GraphIt: a High-Performance Graph DSL
[pdf] [arxiv] [GraphIt project]
Yunming Zhang, Mengjiao Yang, Riyadh Baghdadi, Shoaib Kamil, Julian Shun, and Saman Amarasinghe
Proceedings of the ACM on Programming Languages, Volume 2. (OOPSLA 2018)
Making Caches Work for Graph Analytics
[pdf] [arxiv]
Yunming Zhang, Vladimir Kiriansky, Charith Mendis, Matei Zaharia, Saman Amarasinghe
Proceedings of the 2017 IEEE International Conference on Big Data (Big Data 2017)
Best Student Paper Award
Optimizing Indirect Memory References with Milk
[pdf]
Vladimir Kiriansky, Yunming Zhang, Saman Amarasinghe
Proceedings of the 2016 International Conference on Parallel Architectures and Compilation Techniques (PACT 2016)
Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code
[pdf] [arxiv] [project]
Romdhane, Emanuele Del Sozzo, Abdurrahman Akkas, Yunming Zhang, Patricia Suriana, Shoaib Kamil, Saman Amarasinghe
Proceedings of the 2019 International Symposium on Code Generation and Optimization (CGO 2019)
HJ-Hadoop: an Optimized MapReduce Runtime for Multi-Core Systems
[pdf] [master thesis]
Yunming Zhang, Alan Cox, Vivek Sarkar
Proceedings of the 2013 companion publication for conference on Systems, programming, & applications: software for humanity (SPLASH 2013)

Teaching Experiences

I have enjoyed working as a TA for Performance Engineering of Software Systems at MIT(6.172), Fundamentals of Parallel Computing (COMP322), Advanced Object Oriented Computing (COMP310), and Computational Thinking (COMP 140) at Rice University.