PCC: An End-to-End Compilation Framework for Neural Networks on Photonic-Electronic Accelerators

Photonic computing, known for its high bandwidth and energy efficiency, harnesses physical phenomena in the optical domain to accelerate a wide range of computational operations such as dot product, matrix multiplication, Fourier transform, 1D convolution, and more. However, the multitude of computational operations mentioned above poses challenges in mapping realistic neural network workloads onto underlying photonic hardware. This complexity requires extensive expertise and laborious programming, impeding the practical adoption and deployment of photonic acceleration. To address this gap, we propose an end-to-end compilation framework comprising a Photonic Compiler Collection (PCC). This framework automates the mapping of high-level deep neural network (DNN) specifications onto target architectures of photonic-electronic accelerators. Additionally, we present a method to streamline neural network workloads by leveraging the multilevel intermediate representation (MLIR) and compiler optimization techniques, targeting photonic-specific patterns. Moreover, we conduct a comprehensive case study illustrating the integration of a typical computational operator, the Mach-Zehnder Interferometer (MZI) mesh, into PCC. Our experimental results demonstrate that PCC achieves up to a 4x speedup on DNN workloads compared to handcrafted implementations. In summary, our proposed framework offers a practical and automated solution for compiling, optimizing, and flexibly sup-porting newer operators of photonic devices. We anticipate that our framework will significantly accelerate the development and deployment of photonic applications in real-world AI scenarios.