Stephen Chou

Senior Software Engineer

@ Google

I am passionate about designing computer languages and building compilers that make it easier for programmers to write high-performance software. I have a decade of experience working on a wide variety of production-grade and research compilers. I have also published at multiple top programming language and compiler conferences.

Currently, I am a senior software engineer at Google, where I work on the XLA TPU compiler and help drive bring-ups of the compiler for new generations of TPUs. Before that, I completed my PhD in computer science at MIT under the supervision of Prof. Saman Amarasinghe. I was a primary contributor to the TACO sparse tensor algebra compiler, and I developed new techniques to automatically generate efficient code for computing on sparse tensors that are stored in disparate data layouts.

Projects

TACO

A compiler that generates fast code to perform sparse tensor algebra computations. I extended TACO so that it can emit code to efficiently compute with sparse tensors that are stored in a wide range of specialized data structures. I also helped maintain many other parts of the compiler, as well as developed a web interface that lets users try out TACO without having to install it on their own machines.

simit

A domain-specific programming language for computing on sparse systems using linear algebra. I implemented a new front end for the Simit compiler that was significantly more user-friendly and robust than the original one. This new front end was later also adapted for use in the compiler for GraphIt, a programming language for graph computations.

Publications

Conference Publications

Compilation of Dynamic Sparse Tensor Algebra
Stephen Chou, Saman Amarasinghe
Proceedings of the ACM on Programming Languages, Volume 6, Issue OOPSLA2, 2022.

PDF

Unified Compilation for Lossless Compression and Sparse Computing
Daniel Donenfeld, Stephen Chou, Saman Amarasinghe
Proceedings of the 2022 IEEE/ACM International Symposium on Code Generation and Optimization, 2022.

PDF

Compilation of Sparse Array Programming Models
Rawn Henry, Olivia Hsu, Rohan Yadav, Stephen Chou, Kunle Olukotun, Saman Amarasinghe, Fredrik Kjolstad
Proceedings of the ACM on Programming Languages, Volume 5, Issue OOPSLA, 2021.

PDF

A Sparse Iteration Space Transformation Framework for Sparse Tensor Algebra
Ryan Senanayake, Changwon Hong, Ziheng Wang, Amalee Wilson, Stephen Chou, Shoaib Kamil, Saman Amarasinghe, Fredrik Kjolstad
Proceedings of the ACM on Programming Languages, Volume 4, Issue OOPSLA, 2020.

PDF

Automatic Generation of Efficient Sparse Tensor Format Conversion Routines
Stephen Chou, Fredrik Kjolstad, Saman Amarasinghe
41st ACM SIGPLAN International Conference on Programming Language Design and Implementation (PLDI), 2020.

PDF Slides Video

Format Abstraction for Sparse Tensor Algebra Compilers
Stephen Chou, Fredrik Kjolstad, Saman Amarasinghe
Proceedings of the ACM on Programming Languages, Volume 2, Issue OOPSLA, 2018.

PDF Slides Video

The Tensor Algebra Compiler
Fredrik Kjolstad, Shoaib Kamil, Stephen Chou, David Lugato, Saman Amarasinghe
Proceedings of the ACM on Programming Languages, Volume 1, Issue OOPSLA, 2017.
Distinguished Paper Award

PDF

Peer-Reviewed Short/Workshop Publications

Sparse Tensor Transpositions
Suzanne Mueller, Peter Ahrens, Stephen Chou, Fredrik Kjolstad, Saman Amarasinghe
32nd ACM Symposium on Parallelism in Algorithms and Architectures, 2020.
(brief announcement)

PDF

Taco: compilation et génération de code d’expressions tensorielles
David Lugato, Fredrik Kjolstad, Stephen Chou, Saman Amarasinghe, Shoaib Kamil
Chocs Avancées, No. 12, 2018.

PDF

taco: A Tool to Generate Tensor Algebra Kernels
Fredrik Kjolstad, Stephen Chou, David Lugato, Shoaib Kamil, Saman Amarasinghe
32nd IEEE/ACM International Conference on Automated Software Engineering (ASE), 2017.
(tools paper and demo)

PDF Video

Using Similarity Distance for Performance Prediction of the Query Optimization Process
Anisoara Nica, Stephen Chou
Sixth International Workshop on Testing Database Systems (DBTest), 2013.

PDF

Other Publications

Format Abstractions for the Compilation of Sparse Tensor Algebra
PhD Thesis, Massachusetts Institute of Technology, 2022.

Unified Sparse Formats for Tensor Algebra Compilers
SM Thesis, Massachusetts Institute of Technology, 2018.

Patents

Resource Estimation For A Query Optimization Process
Anisoara Nica, Stephen Chou
Application number: US 13/754,596, Publication number: US9298771 B2, 2013.

Professional Service

Artifact evaluation committee member for CGO 2020
Reviewer for IEEE Transactions on Computers (2019)
Subreviewer for CGO 2017, CGO 2018, PLDI 2019, SPAA 2019, Euro-Par 2020, SPAA 2021, CGO 2022, and MLSys 2022

Awards & Honors

Distinguished Paper Award (OOPSLA, 2017)
Albert Sherwood Barber Medal for Best Overall Work Term and Academic Performance (University of Waterloo, 2015)
First in Class Engineering Scholarship for 4A term (University of Waterloo, 2015)
First in Class Engineering Scholarship for 3B term (University of Waterloo, 2014)
First in Class Engineering Scholarship for 3A term (University of Waterloo, 2013)
iAnywhere Solutions Inc. Scholarship (University of Waterloo, 2012)