UNet Performance with Wafer Scale Engine (Optimization Case Study)

Science-based simulations can significantly improve efficiency and effectiveness of the carbon storage operations and subsurface reservoir management and help to mitigate the impact of atmospheric carbon dioxide on climate change. However, they can be computationally very expensive, especially when inversion modeling is employed. Training robust, trustworthy, and reliable site-specific, field-scale, science-informed artificial intelligence models is also computationally and data intensive. In this work, compute acceleration of such model training was achieved by taking advantage of extreme workflow parallelism of single-wafer-embedding of the compute resources. Successful kernel mapping to system fabric and accelerated training were accomplished using a prototype UNet architecture and dynamic loss scaling.