Harnessing Extreme Heterogeneity for Ocean Modeling with Tensors

Specialized processors designed to accelerate tensor operations are evolving faster than conventional processors. This trend of architectural innovations greatly benefits artificial intelligence (AI) workloads. However, it is unknown how well AI-optimized accelerators can be retargeted to scientific applications. To answer this question we explore (1) whether a typical scientific modeling kernel can be mapped efficiently to tensor operations and (2) whether this approach is portable across diverse processors and AI accelerators. In this paper we implement two versions of tracer advection in an ocean-modeling application using PyTorch and evaluate these on one CPU, two GPUs, and Google's TPU. Our findings are that scientific modeling can observe both a performance boost and improved portability by mapping key computational kernels to tensor operations.