Thread block restructuring and LocalGeometry caching

[tapios]

This draft PR attempts to improve performance in 2 relatively simple ways:

• Restructure thread blocks to process multiple columns per block (all columns in an element), to increase GPU occupancy
• Cache the LocalGeometry (invJ) in shared memory (to have one memory load per column)

@imreddyTeja @dennisYatunin @petebachant Please take a look if any of this is useful for performance improvement, e.g., by looking at an nsys profile.

I won't be able to work on this more right now, so if there's anything useful here, please just take it over.

[petebachant]

Just kicked off a few flagship AMIP jobs to test this out and will report back here.

[tapios]

Thank you!

[petebachant]

Ran this twice with the flagship AMIP 16 helem config. First run was ~1% improvement in SYPD. Second was 0%. Both NSYS runs showed a slowdown for this change compared to baseline [results].

Side note: If we don't have a CI check here that is a solid indicator of flagship AMIP SYPD, we should add one to close the loop for PRs so we don't need to test manually. Might require untangling some circular dependency issues though. There's a nightly run for that config in the Coupler repo, but it uses important dependencies (e.g., ClimaCore) at main and has been broken/disabled for a few days.

[tapios]

Ran this twice with the flagship AMIP 16 helem config. First run was ~1% improvement in SYPD. Second was 0%. Both NSYS runs showed a slowdown for this change compared to baseline [results].

Side note: If we don't have a CI check here that is a solid indicator of flagship AMIP SYPD, we should add one to close the loop for PRs so we don't need to test manually. Might require untangling some circular dependency issues though. There's a nightly run for that config in the Coupler repo, but it uses important dependencies (e.g., ClimaCore) at main and has been broken/disabled for a few days.

Thanks, @petebachant. Really helpful to know. And I agree with your comment on automated CI checks.

I'm curious why processing multiple columns per block shows no benefit. It may be worth looking into this as it may tell us something about structural issues (unless you all already know why that is).

[tapios]

Ran this twice with the flagship AMIP 16 helem config. First run was ~1% improvement in SYPD. Second was 0%. Both NSYS runs showed a slowdown for this change compared to baseline [results].
Side note: If we don't have a CI check here that is a solid indicator of flagship AMIP SYPD, we should add one to close the loop for PRs so we don't need to test manually. Might require untangling some circular dependency issues though. There's a nightly run for that config in the Coupler repo, but it uses important dependencies (e.g., ClimaCore) at main and has been broken/disabled for a few days.

Thanks, @petebachant. Really helpful to know. And I agree with your comment on automated CI checks.

I'm curious why processing multiple columns per block shows no benefit. It may be worth looking into this as it may tell us something about structural issues (unless you all already know why that is).

I see at least one issue: By default, the shared memory branch is disabled:

In finitedifference.jl:
Operators.use_fd_shmem() = false

So the memory optimizations probably were never actually used; the code used the global-memory kernel, and hence there was no difference. We'd need to set ClimaCore.Operators.use_fd_shmem() = true (and then experiment with the number of threads for optimization and to avoid exceed max shmem).