Condenser: Noun, an Apparatus for Compiling Science to the Cloud

Albert Cohen, Research Scientist, Google DeepMind, France

Advances in ML acceleration hardware offer unprecedented opportunities to scale up scientific computing. What about large scale simulations on specialized tensor accelerators with state-of-the-art interconnect technology, also complementing traditional numerical methods with ML modeling? This is where condensers come in: no more vaporware of science in the cloud, but real-world automatic parallelization through domain-specific compilers. Unfortunately scientific applications are built for traditional HPC systems, often written in Fortran, C++ or more recently Julia, and remain largely incompatible with these technologies. Portable abstractions exist, such as Kokkos in C++, but remain much lower level than the compute graphs of popular ML frameworks. The abstraction and domain-specialization gap isolates scientific computing from rapid cloud-based innovation for AI workloads. With a compiler-centric focus, we will showcase recent achievements porting a Julia-based ocean model to both GPU- and Google TPU-based hardware, leveraging the MLIR infrastructure. We will highlight scientific and engineering challenges met along the way, as well as our original path to leverage cutting-edge ML infrastructure, from collective communications to low-level code generation and automatic differentiation.