Zhenyu Bai

My research focuses on dataflow architectures and compiler systems. On the hardware side, we design CGRA-style dataflow architectures, from custom microarchitectures and tapeouts to end-to-end software stacks. On the compiler side, we develop compilation techniques for systems with explicit data movement, distributed memories, and parallel compute units.

A central goal of our work is to support tile-based DSLs such as Helion across both our in-house architectures and commercial platforms, including Tenstorrent, IBM AIU, AMD NPU/AIE, NPU/TPU-style architectures, as well as the new dataflow features of NVidia’s GPUs. Our Loom framework provides an end-to-end compilation flow (Helion → MLIR → TT-Metal) for Tenstorrent, achieving performance comparable to vendor libraries on tensor kernels and fused AI operators.

some of my interesting papers...

ASPLOS’26

A data-driven dynamic execution orchestration architecture

Zhenyu Bai, Pranav Dangi, Rohan Juneja, and 4 more authors

In Proceedings of the 31st ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1, 2026

Bib

@inproceedings{bai2026data,
  title = {A data-driven dynamic execution orchestration architecture},
  author = {Bai, Zhenyu and Dangi, Pranav and Juneja, Rohan and Li, Zhaoying and Yan, Zhanglu and Lan, Huiying and Mitra, Tulika},
  booktitle = {Proceedings of the 31st ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1},
  pages = {1--19},
  year = {2026},
}

OSDI’26

TileLoom: Automatic Dataflow Planning for Tile-Based Languages on Spatial Dataflow Accelerators

Wei Li, Zhenyu Bai, Heru Wang, and 2 more authors

In Symposium on Operating Systems Design and Implementation, 2026

Bib

@inproceedings{li2026tileloom,
  title = {TileLoom: Automatic Dataflow Planning for Tile-Based Languages on Spatial Dataflow Accelerators},
  author = {Li, Wei and Bai, Zhenyu and Wang, Heru and Dangi, Pranav and Zhang, Zhiqiang},
  booktitle = {Symposium on Operating Systems Design and Implementation},
  year = {2026}
}

DAC’24

SWAT: Scalable and efficient window attention-based transformers acceleration on fpgas

Zhenyu Bai, Pranav Dangi, Huize Li, and 1 more author

In Proceedings of the 61st ACM/IEEE Design Automation Conference, 2024

Bib

@inproceedings{bai2024swat,
  title = {SWAT: Scalable and efficient window attention-based transformers acceleration on fpgas},
  author = {Bai, Zhenyu and Dangi, Pranav and Li, Huize and Mitra, Tulika},
  booktitle = {Proceedings of the 61st ACM/IEEE Design Automation Conference},
  pages = {1--6},
  year = {2024},
}

TCAD’26

SparrowSNN: A Hardware/software Co-design for Energy Efficient Low Power Application

Zhenyu Bai, Zhanglu Yan, Bin Gao, and 2 more authors

IEEE Transactions on Computer-Aided Design of Integrated Circuits and Systems, 2026

Bib

@article{bai2026sparrowsnn,
  title = {SparrowSNN: A Hardware/software Co-design for Energy Efficient Low Power Application},
  author = {Bai, Zhenyu and Yan, Zhanglu and Gao, Bin and Mitra, Tulika and Wong, Weng-Fai},
  journal = {IEEE Transactions on Computer-Aided Design of Integrated Circuits and Systems},
  year = {2026},
  publisher = {IEEE},
}

PACT’24

Zed: A generalized accelerator for variably sparse matrix computations in ml

Pranav Dangi, Zhenyu Bai, Rohan Juneja, and 2 more authors

In Proceedings of the 2024 International Conference on Parallel Architectures and Compilation Techniques, 2024

Bib

@inproceedings{dangi2024zed,
  title = {Zed: A generalized accelerator for variably sparse matrix computations in ml},
  author = {Dangi, Pranav and Bai, Zhenyu and Juneja, Rohan and Wijerathne, Dhananjaya and Mitra, Tulika},
  booktitle = {Proceedings of the 2024 International Conference on Parallel Architectures and Compilation Techniques},
  pages = {246--257},
  year = {2024},
}

ASPLOS’25

Enhancing CGRA efficiency through aligned compute and communication provisioning

Zhaoying Li, Pranav Dangi, Chenyang Yin, and 5 more authors

In Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1, 2025

Bib

@inproceedings{li2025enhancing,
  title = {Enhancing CGRA efficiency through aligned compute and communication provisioning},
  author = {Li, Zhaoying and Dangi, Pranav and Yin, Chenyang and Bandara, Thilini Kaushalya and Juneja, Rohan and Tan, Cheng and Bai, Zhenyu and Mitra, Tulika},
  booktitle = {Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1},
  pages = {410--425},
  year = {2025},
}