A GPU compatible version for non-orographic gravity wave parameterization

[Xiaoyang-Xie]

Purpose

Let Nogw parameterization work on GPU

To-do

Speed it up (Now this version now runs 50% slower than the previous version of the code.)

Content

Replace the for-end loop structure with unrolled functions

---

• I have read and checked the items on the review checklist.

[Sbozzolo]

This code is really complex, has little to no comments that explain what is happening, uses advanced features and introduces a dependency on UnrolledUtilities, a package that is not documented in such a way that people can understand what it does and when/how to use it without in-person guidance. I am worried that almost no one will be able to extend and maintain this piece of code.

@dennisYatunin can you help reducing complexity and improving clarity?

If you really want to use UnrolledUtilities, can you first add documentation so that people can understand what the package does and how to use it? If you want to use UnrolledUtilities in ClimaAtmos, the target demographics for your documentation should be graduate students / research scientists.

[szy21]

Suggested change

	dt_save_state_to_disk: "400secs"
	initial_condition: "IsothermalProfile"
	t_end: "1500secs"
	hyperdiff: false
	t_end: "1500secs"

[charleskawczynski]

This overrides a method in ClimaCore. If this is needed, it needs to be put into ClimaCore.

[Xiaoyang-Xie]

Yes, it is needed. Otherwise the code could not run on GPU. Dennis tells me he will open a pull request in ClimaCore soon.

[charleskawczynski]

Can we not hard-code Float32 here? (and below)

[Xiaoyang-Xie]

Sure! I will change it.

[charleskawczynski]

I'm not keen on nested anonymous functions like this. Can we somehow break the contents of this giant function into smaller function calls?

[Xiaoyang-Xie]

I've refactored part of the wave source calculation into smaller function calls. Does this approach work for you? Breaking down the unrolled_reduce function, which is a big part of the whole function, isn’t easy because it directly uses many externally passed parameters. If we avoid using anonymous functions, we’ll have to keep passing those parameters in the reduce operation, which might make the function less elegant.

[charleskawczynski]

wave_source looks great, thanks for updating that. It would be nice if it's possible to do something similar with the block inside the unrolled_reduce(Val(nc), (mask, FT1(0.0))) do (mask, fm), (n) block, but I understand if it's not possible.