# Global Ocean Modeling With GPU Acceleration in Python **URL:** \ **Category:** Uncategorized\ **Created:** [January 5, 2022, 2:59am UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497 "2022-01-05T02:59:02Z")\ **Posts on this page:** 20\ **Page:** 1
**Author:** ![Beliavsky](https://avatars.discourse-cdn.com/v4/letter/b/ba8739/32.png) [@Beliavsky](https://fortran-lang.discourse.group/u/Beliavsky)\ **Post date:** [January 5, 2022, 2:59am UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/1 "2022-01-05T02:59:02Z")
### [Fast, Cheap, and Turbulent—Global Ocean Modeling With GPU Acceleration in Python](https://agupubs.onlinelibrary.wiley.com/doi/10.1029/2021MS002717) by [Dion Häfner](https://agupubs.onlinelibrary.wiley.com/action/doSearch?ContribAuthorRaw=H%C3%A4fner%2C+Dion), [Roman Nuterman](https://agupubs.onlinelibrary.wiley.com/action/doSearch?ContribAuthorRaw=Nuterman%2C+Roman),and [Markus Jochum](https://agupubs.onlinelibrary.wiley.com/action/doSearch?ContribAuthorRaw=Jochum%2C+Markus) Journal of Advances in Modeling Earth Systems First published: 02 December 2021 ### Abstract Even to this date, most earth system models are coded in Fortran, especially those used at the largest compute scales. Our ocean model Veros takes a different approach: it is implemented using the high-level programming language Python. Besides numerous usability advantages, this allows us to leverage modern high-performance frameworks that emerged in tandem with the machine learning boom. By interfacing with the JAX library, Veros is able to run high-performance simulations on both central processing units (CPU) and graphical processing unit (GPU) through the same code base, with full support for distributed architectures. On CPU, Veros is able to match the performance of a Fortran reference, both on a single process and on hundreds of CPU cores. On GPU, we find that each device can replace dozens to hundreds of CPU cores, at a fraction of the energy consumption. We demonstrate the viability of using GPUs for earth system modeling by integrating a global 0.1° eddy-resolving setup in single precision, where we achieve 1.3 model years per day on a single compute instance with 16 GPUs, comparable to over 2,000 Fortran processes. Related [blog post](https://dionhaefner.github.io/2021/12/supercharged-high-resolution-ocean-simulation-with-jax/) ---
**Author:** ![zaikunzhang](https://yyz2.discourse-cdn.com/free1/user_avatar/fortran-lang.discourse.group/zaikunzhang/32/785_2.png) [@zaikunzhang](https://fortran-lang.discourse.group/u/zaikunzhang)\ **Post date:** [January 5, 2022, 3:11am UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/2 "2022-01-05T03:11:02Z")
Unfortunately, Fortran is deemed here as a synonym of “old-fashioned” and “to-be-replaced” (again). What is the remedy for this widespread misconception? ---
**Author:** ![Beliavsky](https://avatars.discourse-cdn.com/v4/letter/b/ba8739/32.png) [@Beliavsky](https://fortran-lang.discourse.group/u/Beliavsky)\ **Post date:** [January 5, 2022, 3:17am UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/3 "2022-01-05T03:17:18Z")
The authors did not make general criticisms of Fortran but have asserted that on the hardware available to them, a Python/JAX solution was faster. Based on this paper, it needs to be easier for Fortran to exploit GPUs. ---
**Author:** ![certik](https://yyz2.discourse-cdn.com/free1/user_avatar/fortran-lang.discourse.group/certik/32/4_2.png) [@certik](https://fortran-lang.discourse.group/u/certik)\ **Post date:** [January 5, 2022, 4:32am UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/4 "2022-01-05T04:32:54Z")
They use the [GitHub - google/jax: Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more](https://github.com/google/jax) library to compile Python to very fast native code using [XLA](https://www.tensorflow.org/xla). If this can optimize array code very well, we could write an XLA backend to LFortran also. ---
**Author:** ![shahmoradi](https://yyz2.discourse-cdn.com/free1/user_avatar/fortran-lang.discourse.group/shahmoradi/32/3151_2.png) [@shahmoradi](https://fortran-lang.discourse.group/u/shahmoradi)\ **Post date:** [January 5, 2022, 4:56am UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/5 "2022-01-05T04:56:36Z")
This is the modernized equivalent of their Fortran code in Figure 2, ```fortran do concurrent(j=js_pe:je_pe, i=is_pe:ie_pe, tke(i,j,nz,taup1) < 0.) tke_surf_corr(i,j) = -tke(i,j,nz,taup1) * 0.5 * dzw(nz) / dt_tke tke(i,j,nz,taup1) = 0. end do ``` This code does not even need extra jit decorators to run on GPUs. Compile it with NVIDIA compilers to get GPU parallelism. want CPU parallelism? compile and run instead with Intel ifort. Just compare the elegance of this code with the Python equivalent in their figure. The authors state further “When going from Fortran to Python, explicit loops are replaced with array operations”. Why? Is it not possible in Fortran? The entire vectorization ideas of numpy are inherited from Fortran 90. ---
**Author:** ![dionhaefner](https://avatars.discourse-cdn.com/v4/letter/d/b487fb/32.png) [@dionhaefner](https://fortran-lang.discourse.group/u/dionhaefner)\ **Post date:** [January 5, 2022, 9:53am UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/6 "2022-01-05T09:53:24Z")
👋 Hey everybody, author of the article here. Happy to discuss this with you. To address some of your points: > [@Beliavsky](#): > > Based on this paper, it needs to be easier for Fortran to exploit GPUs. Yes! Unfortunately, hardware agnostic code doesn’t mix very well with explicit loops, so unless you restrict yourself to vectorized operations I don’t have much hope for Fortran here (would love to be wrong though). > [@certik](#): > > If this can optimize array code very well, we could write an XLA backend to LFortran also That would be amazing, but see the caveat above: XLA can _only_ accelerate array code, so you would have to stick to a subset of Fortran (or do some serious compiler magic). > [@shahmoradi](#): > > This is the modernized equivalent of their Fortran code in Figure 2, > > ```auto > do concurrent(j=js_pe:je_pe, i=is_pe:ie_pe, tke(i,j,nz,taup1) < 0.) > tke_surf_corr(i,j) = -tke(i,j,nz,taup1) * 0.5 * dzw(nz) / dt_tke > tke(i,j,nz,taup1) = 0. > end do > > ``` That looks great and very clean! I wouldn’t say that Python code is inherently more readable than Fortran code, and the inverse can often be true. On the other hand, Python gives you all the tools you need to build your own abstractions, so once you have a working + fast implementation you can start iterating towards beauty (this is next up on our roadmap). I discuss this in a different blog post [here](https://dionhaefner.github.io/2021/04/higher-level-geophysical-modelling/#abstraction-is-key). > [@shahmoradi](#): > > This code does not even need extra jit decorators to run on GPUs. Compile it with NVIDIA compilers to get GPU parallelism. want CPU parallelism? compile and run instead with Intel ifort. Just compare the elegance of this code with the Python equivalent in their figure. This sounds good on paper but is not (yet) flexible enough for complex applications like ours (see [my response to your comment](https://github.com/dionhaefner/blog-comments/issues/4#issuecomment-1005523361) on the article). Also, we’re typically not interested in thread-based parallelism since we run on multiple nodes anyway. > [@shahmoradi](#): > > The authors state further “When going from Fortran to Python, explicit loops are replaced with array operations”. Why? Is it not possible in Fortran? The entire vectorization ideas of numpy are inherited from Fortran 90. This is out of necessity, not virtue. ---
**Author:** ![ivanpribec](https://yyz2.discourse-cdn.com/free1/user_avatar/fortran-lang.discourse.group/ivanpribec/32/3290_2.png) [@ivanpribec](https://fortran-lang.discourse.group/u/ivanpribec)\ **Post date:** [January 5, 2022, 10:49am UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/7 "2022-01-05T10:49:53Z")
> [@shahmoradi](#): > > ```fortran > do concurrent(j=js_pe:je_pe, i=is_pe:ie_pe, tke(i,j,nz,taup1) < 0.) > tke_surf_corr(i,j) = -tke(i,j,nz,taup1) * 0.5 * dzw(nz) / dt_tke > tke(i,j,nz,taup1) = 0. > end do > > ``` This actually feels cleaner to me, particularly the last line, compared to: ```python tke = tke.at[2:-2,2:-2,-1,taup1].multiply(mask) ``` in the Python + JAX version. It took me half a minute to understand that since the mask is an array of zeros and ones, the elements multipled by 1 will be left “untouched”, while the rest will be set to zero. ---
**Author:** ![milancurcic](https://yyz2.discourse-cdn.com/free1/user_avatar/fortran-lang.discourse.group/milancurcic/32/2_2.png) [@milancurcic](https://fortran-lang.discourse.group/u/milancurcic)\ **Post date:** [January 5, 2022, 12:28pm UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/8 "2022-01-05T12:28:39Z")
Welcome @dionhaefner, and thank you for posting here! ---
**Author:** ![ivanpribec](https://yyz2.discourse-cdn.com/free1/user_avatar/fortran-lang.discourse.group/ivanpribec/32/3290_2.png) [@ivanpribec](https://fortran-lang.discourse.group/u/ivanpribec)\ **Post date:** [January 5, 2022, 2:07pm UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/9 "2022-01-05T14:07:19Z")
The current situation with Python libraries for multi-dimensional arrays including NumPy, Tensorflow, PyTorch, Dask, JAX, CuPy, MXNet, Xarray and others, reminds me strongly of early Fortran implementations. The inconsistencies between these Python libraries has led to the launch of a [Python Array API Standard](https://data-apis.org/array-api/latest/purpose_and_scope.html). While Python is often applauded for the high productivity of programmers, it seems to suffer the same amount of “reinventing the wheel” like any other programming language. I’ve read your [blog post](https://dionhaefner.github.io/2021/04/higher-level-geophysical-modelling/#comments) and agree with most of the takeaways. One nitpick I have is the term “modern Fortran”, which is actually fairly old and has been in use all the way back since Fortran 90 if not earlier (I suppose that’s why you added the quotation marks). The abstraction “separate numerics from physics” is something which I’ve encountered before in the book of Damian Rouson (@rouson) - [Scientific Software Design: The Object-Oriented Way](https://books.google.de/books/about/Scientific_Software_Design.html?id=xedE5KkHVn4C&source=kp_book_description&redir_esc=y). Some authors also call this “coordinate-free programming”; here are just a few papers on the topic: - [Coordinate Free Programming of Computational Fluid Dynamics Problems](https://www.hindawi.com/journals/sp/2000/419840/) - [Use of coordinate-free numerics in elastic wave simulation - ScienceDirect](https://www.sciencedirect.com/science/article/abs/pii/S0168927401000964) - [Coordinate Free Numerics Part 1 - How to avoid index wrestling in tensor computations](http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.17.7672&rep=rep1&type=pdf) - [Coordinate-free numerics: all your variation points for free? | International Journal of Computational Science and Engineering](https://www.inderscienceonline.com/doi/abs/10.1504/IJCSE.2009.029159) - [On Object-Oriented Frameworks and Coordinate Free Formulations of PDEs | Engineering with Computers](https://link.springer.com/content/pdf/10.1007/s003660200026.pdf) ---
**Author:** ![certik](https://yyz2.discourse-cdn.com/free1/user_avatar/fortran-lang.discourse.group/certik/32/4_2.png) [@certik](https://fortran-lang.discourse.group/u/certik)\ **Post date:** [January 5, 2022, 3:31pm UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/10 "2022-01-05T15:31:26Z")
First of all, thanks @dionhaefner for writing here! > [@dionhaefner](#): > > That would be amazing, but see the caveat above: XLA can _only_ accelerate array code, so you would have to stick to a subset of Fortran (or do some serious compiler magic). Would you have time for a quick video call to discuss this further? I would love to learn from your experience. What I had in mind is to combine LLVM with XLA, or whatever else is needed to get it running. I don’t see fundamentally how LFortran is different from JAX: the code from your article that uses the `@jax.jit` decorator could be equivalently written in Fortran and then LFortran can do whatever JAX is doing, in principle. However, the devil is in the details and I would love to hear your experience and brainstorm this more. I completely agree with you that Fortran has ways to go in order to be as usable as Python, and that is precisely why we are trying to improve the whole Fortran ecosystem. As a language Fortran seems pretty nice and I don’t see any hard roadblock, it’s “just” that the compilers must improve and the other tooling, and that is exactly what we are working on. ---
**Author:** ![dionhaefner](https://avatars.discourse-cdn.com/v4/letter/d/b487fb/32.png) [@dionhaefner](https://fortran-lang.discourse.group/u/dionhaefner)\ **Post date:** [January 5, 2022, 3:35pm UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/11 "2022-01-05T15:35:06Z")
> [@certik](#): > > Would you have time for a quick video call to discuss this further? I would love to learn from your experience. Sure. I couldn’t figure out how to send a private message here, but if you just send me an e-mail at dion.haefner@nbi.ku.dk we can bang something out. > [@ivanpribec](#): > > While Python is often applauded for the high productivity of programmers, it seems to suffer the same amount of “reinventing the wheel” like any other programming language. I agree, and I hope we’ll have those community standards for accelerator frameworks soon. > [@ivanpribec](#): > > The abstraction “separate numerics from physics” is something which I’ve encountered before in the book of Damian Rouson (@rouson) - [Scientific Software Design: The Object-Oriented Way](https://books.google.de/books/about/Scientific_Software_Design.html?id=xedE5KkHVn4C&source=kp_book_description&redir_esc=y). Some authors also call this “coordinate-free programming”; here are just a few papers on the topic: Thanks for the pointers! These will come in handy when we get deeper into this discussion. I also quite like [the abstractions used by jax-cfd](https://github.com/google/jax-cfd/blob/main/jax_cfd/base/diffusion.py), but if you take it too far you’re tossing code reuse out the window, so there’s always a trade-off. ---
**Author:** ![certik](https://yyz2.discourse-cdn.com/free1/user_avatar/fortran-lang.discourse.group/certik/32/4_2.png) [@certik](https://fortran-lang.discourse.group/u/certik)\ **Post date:** [January 5, 2022, 4:02pm UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/12 "2022-01-05T16:02:47Z")
> [@dionhaefner](#): > > Sure. I couldn’t figure out how to send a private message here, but if you just send me an e-mail at dion.haefner@nbi.ku.dk we can bang something out. Thanks! I sent you an email. ---
**Author:** ![shahmoradi](https://yyz2.discourse-cdn.com/free1/user_avatar/fortran-lang.discourse.group/shahmoradi/32/3151_2.png) [@shahmoradi](https://fortran-lang.discourse.group/u/shahmoradi)\ **Post date:** [January 6, 2022, 11:15pm UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/13 "2022-01-06T23:15:21Z")
Thanks for the clarifications. I have not been able to read your GitHub response yet. I wish the new Fortran features were also discussed in this article. Now, people who read the paper may (implicitly) get the (incorrect) impression that Fortran fundamentally lacks cross-platform syntax. ---
**Author:** ![dionhaefner](https://avatars.discourse-cdn.com/v4/letter/d/b487fb/32.png) [@dionhaefner](https://fortran-lang.discourse.group/u/dionhaefner)\ **Post date:** [January 7, 2022, 8:40am UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/14 "2022-01-07T08:40:38Z")
I would have mentioned the NVIDIA SDK if I had known about it, but that doesn’t change the core message. > [@shahmoradi](#): > > Now, people who read the paper may (implicitly) get the (incorrect) impression that Fortran fundamentally lacks cross-platform syntax. Fortran _does_ lack cross-platform syntax. Running your code through a vendor-specific compiler that can only handle the simplest cases of embarassingly parallel workloads is _not_ the same as having access to a fully integrated GPU framework. I don’t think denying the problems of the language / ecosystem does anyone a favor. I’m the first one to admit that Python has many problems that are not present in other languages (including Fortran), but in this case the reverse is true, and the earlier we agree on that the earlier we can change it. ---
**Author:** ![shahmoradi](https://yyz2.discourse-cdn.com/free1/user_avatar/fortran-lang.discourse.group/shahmoradi/32/3151_2.png) [@shahmoradi](https://fortran-lang.discourse.group/u/shahmoradi)\ **Post date:** [January 7, 2022, 9:00am UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/15 "2022-01-07T09:00:59Z")
> [@dionhaefner](#): > > Fortran _does_ lack cross-platform syntax. I’m afraid that’s not right. Please be careful to speak the truth when you speak of an entire language ecosystem (that is all I am asking of the authors of this article too). Fortran **does** have the cross-platform cross-architecture **syntax**. Do you know about Fortran Coarrays? Do you know about the parallelism paradigm it supports? Fortran is the pioneer of all languages in this regard when Coarrays were the only parallelism paradigm three decades ago while no other programming language had any notion of parallelism. A separate issue is whether the compilers can utilize this complex yet simple, beautiful syntax of Fortran parallelism to implement cross-arch parallelism paradigms. Some compilers like Intel ifort and OpenCoarrays library can convert the Fortran syntax to RMA MPI communications for distributed parallelism, some others like NAG convert **the same Coarray** code to shared-memory parallelism, and some like NVIDIA are working to convert **the same code** to GPU parallelism. ---
**Author:** ![dionhaefner](https://avatars.discourse-cdn.com/v4/letter/d/b487fb/32.png) [@dionhaefner](https://fortran-lang.discourse.group/u/dionhaefner)\ **Post date:** [January 7, 2022, 9:52am UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/16 "2022-01-07T09:52:17Z")
> [@shahmoradi](#): > > and some like NVIDIA are working to convert **the same code** to GPU parallelism. If this bears fruit I’m happy to change my mind, but until then, the point stands. ---
**Author:** ![oscardssmith](https://avatars.discourse-cdn.com/v4/letter/o/b9e5f3/32.png) [@oscardssmith](https://fortran-lang.discourse.group/u/oscardssmith)\ **Post date:** [January 7, 2022, 1:26pm UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/17 "2022-01-07T13:26:04Z")
That only works if you only care about 1 type of parallelism. What happens if you want to use multi-threaded code that also uses GPU? For an example, look at AlphaZero/Lc0. They use multi-threaded MCTS search combined with GPU evaluation of neural networks. How would you deal with this in Fortran? ---
**Author:** ![ivanpribec](https://yyz2.discourse-cdn.com/free1/user_avatar/fortran-lang.discourse.group/ivanpribec/32/3290_2.png) [@ivanpribec](https://fortran-lang.discourse.group/u/ivanpribec)\ **Post date:** [January 7, 2022, 2:31pm UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/18 "2022-01-07T14:31:31Z")
I am not familiar with AlphaZero, neural networks, or MCTS, but just wanted to mention [FEATS](https://github.com/sourceryinstitute/FEATS) which is a Fortran 2018 framework for parallel asynchronous task-scheduling. As far as I can understand, the concept of Fortran coarray images is supposed to be agnostic to the type of underlying hardware. In principle you could have a coarray team consisting of multi-threaded CPU’s, and a second coarray team of multi-threaded GPU’s. Granted, this is not yet something we see in practice so from a pragmatic point of view Dion’s point is valid. I’m hoping that the gfortran implementation of coarrays using OpenMP might be able to leverage both types of compute platforms in the future. Similar to the abstraction “separate numerics from the physics”, Fortran’s goal has always been to “separate numerics from the hardware”. Following this line of thought, I am a bit critical of the title of the paper, as it implies that Python and GPU acceleration are somehow needed if you want to do state of the art ocean modelling, even though the equations solved in this paper are practically the same as in the rest of the models. I think there is a crucial distinction between ocean model (which is the set of equations, boundary conditions, and parameters), and ocean model _implementation_ such as Veros. Conflating the two can be a source of misunderstanding. In a few decades the hardware will change again, and we’ll all have the opportunity to write a bunch of new papers “Modeling X with Y Hardware in Z Language”. ---
**Author:** ![R\_cubed](https://yyz2.discourse-cdn.com/free1/user_avatar/fortran-lang.discourse.group/r_cubed/32/263_2.png) [@R\_cubed](https://fortran-lang.discourse.group/u/R_cubed)\ **Post date:** [January 7, 2022, 2:37pm UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/19 "2022-01-07T14:37:06Z")
> [@ivanpribec](#): > > In a few decades the hardware will change again, and we’ll all have the opportunity to write a bunch of new papers “Modeling X with Y Hardware in Z Language”. In publish or perish academia, that sounds like a feature, not a bug. @certik wrote: > Blockquote > They use the GitHub - google/jax: Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more 2 library to compile Python to very fast native code using XLA 2. If this can optimize array code very well, we could write an XLA backend to LFortran also. From XLA website: > Blockquote > XLA provides an alternative mode of running models: it compiles the TensorFlow graph into a sequence of computation kernels generated specifically for the given model. **Because these kernels are unique to the model, they can exploit model-specific information for optimization**. Correct me if I’m wrong, but this seemed like a paper describing a compiler optimization for array algorithms. The only thing keeping Fortran from implementing this is resources. ---
**Author:** ![certik](https://yyz2.discourse-cdn.com/free1/user_avatar/fortran-lang.discourse.group/certik/32/4_2.png) [@certik](https://fortran-lang.discourse.group/u/certik)\ **Post date:** [January 7, 2022, 3:49pm UTC](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497/20 "2022-01-07T15:49:39Z")
> [@R\_cubed](#): > > Correct me if I’m wrong, but this seemed like a paper describing a compiler optimization for array algorithms. The only thing keeping Fortran from implementing this is resources. I think that is correct. I would like LFortran to be able to use XLA as an option. [Next page](https://fortran-lang.discourse.group/t/global-ocean-modeling-with-gpu-acceleration-in-python/2497.md?page=2)