Define a linear partition, and use in FD stencils

[charleskawczynski]

A performance regression was found in #1969. Concretely,

Main

Problem size: (4, 4, 1, 63, 1536), N-reps: 1,  Float_type = Float64, Device_bandwidth_GBs=2039
┌───────────────────────────────────────────────────────────────┬───────────────────────────────────┬───────────┬─────────────┬────────────────┐
│ funcs                                                         │ time per call                     │ bw %      │ achieved bw │ N reads-writes │
├───────────────────────────────────────────────────────────────┼───────────────────────────────────┼───────────┼─────────────┼────────────────┤
│ (op_GradientF2C!, :none)                                      │ 31 microseconds, 280 nanoseconds  │ 36.1744   │ 737.597     │ 2              │
│ (op_GradientF2C!, :SetValue, :SetValue)                       │ 45 microseconds, 911 nanoseconds  │ 24.6461   │ 502.533     │ 2              │
│ (op_GradientC2F!, :SetGradient, :SetGradient)                 │ 49 microseconds, 740 nanoseconds  │ 22.7483   │ 463.838     │ 2              │
│ (op_GradientC2F!, :SetValue, :SetValue)                       │ 41 microseconds, 890 nanoseconds  │ 27.0119   │ 550.772     │ 2              │
│ (op_DivergenceF2C!, :none)                                    │ 2 milliseconds, 101 microseconds  │ 0.807571  │ 16.4664     │ 3              │
│ (op_DivergenceF2C!, :Extrapolate, :Extrapolate)               │ 496 microseconds, 567 nanoseconds │ 3.41798   │ 69.6925     │ 3              │
│ (op_DivergenceC2F!, :SetDivergence, :SetDivergence)           │ 77 microseconds, 619 nanoseconds  │ 21.8664   │ 445.856     │ 3              │
│ (op_InterpolateF2C!, :none)                                   │ 31 microseconds, 831 nanoseconds  │ 35.5482   │ 724.828     │ 2              │
│ (op_InterpolateC2F!, :SetValue, :SetValue)                    │ 43 microseconds, 859 nanoseconds  │ 25.7986   │ 526.033     │ 2              │
│ (op_InterpolateC2F!, :Extrapolate, :Extrapolate)              │ 39 microseconds, 699 nanoseconds  │ 28.502    │ 581.155     │ 2              │
│ (op_broadcast_example0!, :none)                               │ 33 microseconds, 870 nanoseconds  │ 50.1122   │ 1021.79     │ 3              │
│ (op_broadcast_example1!, :none)                               │ 63 microseconds, 279 nanoseconds  │ 35.7623   │ 729.193     │ 4              │
│ (op_broadcast_example2!, :none)                               │ 62 microseconds, 560 nanoseconds  │ 36.1739   │ 737.585     │ 4              │
│ (op_LeftBiasedC2F!, :SetValue)                                │ 33 microseconds, 590 nanoseconds  │ 33.6856   │ 686.85      │ 2              │
│ (op_LeftBiasedF2C!, :none)                                    │ 30 microseconds, 740 nanoseconds  │ 36.8087   │ 750.53      │ 2              │
│ (op_LeftBiasedF2C!, :SetValue)                                │ 32 microseconds, 999 nanoseconds  │ 34.2889   │ 699.151     │ 2              │
│ (op_RightBiasedC2F!, :SetValue)                               │ 33 microseconds, 980 nanoseconds  │ 33.299    │ 678.967     │ 2              │
│ (op_RightBiasedF2C!, :none)                                   │ 30 microseconds, 761 nanoseconds  │ 36.7848   │ 750.042     │ 2              │
│ (op_RightBiasedF2C!, :SetValue)                               │ 33 microseconds, 910 nanoseconds  │ 33.3677   │ 680.368     │ 2              │
│ (op_CurlC2F!, :SetCurl, :SetCurl)                             │ 56 microseconds, 891 nanoseconds  │ -9.94463  │ -202.771    │ -1             │
│ (op_CurlC2F!, :SetValue, :SetValue)                           │ 1 millisecond, 225 microseconds   │ -0.461728 │ -9.41462    │ -1             │
│ (op_UBPC2F!, :SetValue, :SetValue)                            │ 68 microseconds, 790 nanoseconds  │ -8.22443  │ -167.696    │ -1             │
│ (op_UBPC2F!, :Extrapolate, :Extrapolate)                      │ 69 microseconds, 570 nanoseconds  │ -8.13222  │ -165.816    │ -1             │
│ (op_divO3UBPC2F!, :1SidedO3, :1SidedO3, :SetValue, :SetValue) │ 244 microseconds, 58 nanoseconds  │ -2.3181   │ -47.266     │ -1             │
│ (op_divgrad_CC!, :SetValue, :SetValue, :none)                 │ 134 microseconds, 919 nanoseconds │ 12.5798   │ 256.502     │ 3              │
│ (op_divgrad_FF!, :none, :SetDivergence, :SetDivergence)       │ 59 microseconds, 341 nanoseconds  │ 28.6021   │ 583.197     │ 3              │
│ (op_div_interp_CC!, :SetValue, :SetValue, :none)              │ 586 microseconds, 486 nanoseconds │ -0.964645 │ -19.6691    │ -1             │
│ (op_div_interp_FF!, :none, :SetValue, :SetValue)              │ 69 microseconds, 470 nanoseconds  │ -8.14392  │ -166.055    │ -1             │
│ (op_divgrad_uₕ!, :none, :SetValue, :Extrapolate)              │ 1 millisecond, 844 microseconds   │ -0.306743 │ -6.2545     │ -1             │
│ (op_divgrad_uₕ!, :none, :SetValue, :SetValue)                 │ 114 microseconds, 140 nanoseconds │ -4.95668  │ -101.067    │ -1             │
└───────────────────────────────────────────────────────────────┴───────────────────────────────────┴───────────┴─────────────┴────────────────┘

3fd62e1

Problem size: (4, 4, 1, 63, 1536), N-reps: 1,  Float_type = Float64, Device_bandwidth_GBs=2039
┌───────────────────────────────────────────────────────────────┬───────────────────────────────────┬──────────┬─────────────┬────────────────┐
│ funcs                                                         │ time per call                     │ bw %     │ achieved bw │ N reads-writes │
├───────────────────────────────────────────────────────────────┼───────────────────────────────────┼──────────┼─────────────┼────────────────┤
│ (op_GradientF2C!, :none)                                      │ 28 microseconds, 900 nanoseconds  │ 39.1536  │ 798.342     │ 2              │
│ (op_GradientF2C!, :SetValue, :SetValue)                       │ 39 microseconds, 41 nanoseconds   │ 28.9831  │ 590.965     │ 2              │
│ (op_GradientC2F!, :SetGradient, :SetGradient)                 │ 43 microseconds, 701 nanoseconds  │ 25.8925  │ 527.947     │ 2              │
│ (op_GradientC2F!, :SetValue, :SetValue)                       │ 34 microseconds, 959 nanoseconds  │ 32.3665  │ 659.953     │ 2              │
│ (op_DivergenceF2C!, :none)                                    │ 68 microseconds, 559 nanoseconds  │ 24.7561  │ 504.776     │ 3              │
│ (op_DivergenceF2C!, :Extrapolate, :Extrapolate)               │ 88 microseconds, 871 nanoseconds  │ 19.0981  │ 389.411     │ 3              │
│ (op_DivergenceC2F!, :SetDivergence, :SetDivergence)           │ 67 microseconds, 650 nanoseconds  │ 25.0887  │ 511.559     │ 3              │
│ (op_InterpolateF2C!, :none)                                   │ 29 microseconds, 61 nanoseconds   │ 38.9367  │ 793.919     │ 2              │
│ (op_InterpolateC2F!, :SetValue, :SetValue)                    │ 30 microseconds, 320 nanoseconds  │ 37.3186  │ 760.926     │ 2              │
│ (op_InterpolateC2F!, :Extrapolate, :Extrapolate)              │ 33 microseconds, 530 nanoseconds  │ 33.7469  │ 688.1       │ 2              │
│ (op_broadcast_example0!, :none)                               │ 47 microseconds, 409 nanoseconds  │ 35.8002  │ 729.965     │ 3              │
│ (op_broadcast_example1!, :none)                               │ 82 microseconds, 741 nanoseconds  │ 27.3507  │ 557.682     │ 4              │
│ (op_broadcast_example2!, :none)                               │ 82 microseconds, 599 nanoseconds  │ 27.3974  │ 558.634     │ 4              │
│ (op_LeftBiasedC2F!, :SetValue)                                │ 28 microseconds, 440 nanoseconds  │ 39.7869  │ 811.255     │ 2              │
│ (op_LeftBiasedF2C!, :none)                                    │ 28 microseconds, 330 nanoseconds  │ 39.94    │ 814.377     │ 2              │
│ (op_LeftBiasedF2C!, :SetValue)                                │ 29 microseconds, 461 nanoseconds  │ 38.408   │ 783.139     │ 2              │
│ (op_RightBiasedC2F!, :SetValue)                               │ 28 microseconds, 820 nanoseconds  │ 39.2623  │ 800.558     │ 2              │
│ (op_RightBiasedF2C!, :none)                                   │ 27 microseconds, 511 nanoseconds  │ 41.1305  │ 838.651     │ 2              │
│ (op_RightBiasedF2C!, :SetValue)                               │ 30 microseconds, 19 nanoseconds   │ 37.6928  │ 768.556     │ 2              │
│ (op_CurlC2F!, :SetCurl, :SetCurl)                             │ 51 microseconds, 860 nanoseconds  │ -10.9092 │ -222.438    │ -1             │
│ (op_CurlC2F!, :SetValue, :SetValue)                           │ 54 microseconds, 509 nanoseconds  │ -10.379  │ -211.628    │ -1             │
│ (op_UBPC2F!, :SetValue, :SetValue)                            │ 54 microseconds, 660 nanoseconds  │ -10.3503 │ -211.044    │ -1             │
│ (op_UBPC2F!, :Extrapolate, :Extrapolate)                      │ 58 microseconds, 159 nanoseconds  │ -9.72764 │ -198.347    │ -1             │
│ (op_divO3UBPC2F!, :1SidedO3, :1SidedO3, :SetValue, :SetValue) │ 201 microseconds, 219 nanoseconds │ -2.81163 │ -57.3291    │ -1             │
│ (op_divgrad_CC!, :SetValue, :SetValue, :none)                 │ 135 microseconds, 579 nanoseconds │ 12.5185  │ 255.253     │ 3              │
│ (op_divgrad_FF!, :none, :SetDivergence, :SetDivergence)       │ 54 microseconds, 910 nanoseconds  │ 30.9102  │ 630.26      │ 3              │
│ (op_div_interp_CC!, :SetValue, :SetValue, :none)              │ 88 microseconds, 690 nanoseconds  │ -6.37896 │ -130.067    │ -1             │
│ (op_div_interp_FF!, :none, :SetValue, :SetValue)              │ 62 microseconds, 511 nanoseconds  │ -9.05055 │ -184.541    │ -1             │
│ (op_divgrad_uₕ!, :none, :SetValue, :Extrapolate)              │ 105 microseconds, 959 nanoseconds │ -5.33933 │ -108.869    │ -1             │
│ (op_divgrad_uₕ!, :none, :SetValue, :SetValue)                 │ 92 microseconds, 210 nanoseconds  │ -6.13552 │ -125.103    │ -1             │
└───────────────────────────────────────────────────────────────┴───────────────────────────────────┴──────────┴─────────────┴────────────────┘

Most notably:

main

│ (op_DivergenceF2C!, :none)                                    │ 2 milliseconds, 101 microseconds  │ 0.807571  │ 16.4664     │ 3              │
│ (op_DivergenceF2C!, :Extrapolate, :Extrapolate)               │ 496 microseconds, 567 nanoseconds │ 3.41798   │ 69.6925     │ 3              │
│ (op_CurlC2F!, :SetValue, :SetValue)                           │ 1 millisecond, 225 microseconds   │ -0.461728 │ -9.41462    │ -1             │
│ (op_divgrad_uₕ!, :none, :SetValue, :Extrapolate)              │ 1 millisecond, 844 microseconds   │ -0.306743 │ -6.2545     │ -1             │

3fd62e1

│ (op_DivergenceF2C!, :none)                                    │ 68 microseconds, 559 nanoseconds  │ 24.7561  │ 504.776     │ 3              │
│ (op_DivergenceF2C!, :Extrapolate, :Extrapolate)               │ 88 microseconds, 871 nanoseconds  │ 19.0981  │ 389.411     │ 3              │
│ (op_CurlC2F!, :SetValue, :SetValue)                           │ 54 microseconds, 509 nanoseconds  │ -10.379  │ -211.628    │ -1             │
│ (op_divgrad_uₕ!, :none, :SetValue, :Extrapolate)              │ 105 microseconds, 959 nanoseconds │ -5.33933 │ -108.869    │ -1             │

This PR is an attempt to fix this regression by reverting the prescribed thread-block configuration for stencils to use a linear partition.

[charleskawczynski]

@Sbozzolo or @szy21, can one of you please try this out to see if this fixes the performance regression in the coupler job?

[charleskawczynski]

For now, this fully reverts using the multi-dimensional thread-block to using a dynamic CartesianIndices, however, we could try to pass CartesianIndices through Val into the kernel, and see if this offers any speedup (in case the compiler can optimize division by constant divisors).

[Sbozzolo]

@Sbozzolo or @szy21, can one of you please try this out to see if this fixes the performance regression in the coupler job?

Here: https://buildkite.com/clima/climacoupler-amip/builds/56

[charleskawczynski]

@Sbozzolo or @szy21, can one of you please try this out to see if this fixes the performance regression in the coupler job?

Here: https://buildkite.com/clima/climacoupler-amip/builds/56

Looks like it's broken on an unrelated land component.

[Sbozzolo]

@Sbozzolo or @szy21, can one of you please try this out to see if this fixes the performance regression in the coupler job?

Here: https://buildkite.com/clima/climacoupler-amip/builds/56

Looks like it's broken on an unrelated land component.

The update in Thermodynamics broke the coupler tests, so we cannot (trivially) merge the fix for this. I'll merge them manually

[Sbozzolo]

I fixed it, but now it crashed because of the new thermodynamics

[Sbozzolo]

This is running: https://buildkite.com/clima/climacoupler-amip/builds/64#019229f8-e641-4bc4-b733-99e276fb74d9 (at the time of writing, this is the only job running and only GPU being used)

[szy21]

This is running: https://buildkite.com/clima/climacoupler-amip/builds/64#019229f8-e641-4bc4-b733-99e276fb74d9 (at the time of writing, this is the only job running and only GPU being used)

There seems to be a ~2% increase in SYPD in this run compared with the current main. Not sure if it is significant?

[charleskawczynski]

There seems to be a ~2% increase in SYPD in this run compared with the current main. Not sure if it is significant?

I can try reverting more kernel launches, we don't have any microbenchmarks for the TDMA and multiple field solves. That should only account for an additional 4%, though, since Atmos was partially responsible for the slowdown (I believe @szy21 mentioned 8% total, 6% from ClimaCore, and 2% from Atmos). If that fails to recover the last few percent, then we can just revert all of ClimaCore back to that commit and reapply changes since.

[charleskawczynski]

This is running: https://buildkite.com/clima/climacoupler-amip/builds/64#019229f8-e641-4bc4-b733-99e276fb74d9

I'd like to iterate on this more quickly, how can I run this job interactively?

[Sbozzolo]

There seems to be a ~2% increase in SYPD in this run compared with the current main. Not sure if it is significant?

I can try reverting more kernel launches, we don't have any microbenchmarks for the TDMA and multiple field solves. That should only account for an additional 4%, though, since Atmos was partially responsible for the slowdown (I believe @szy21 mentioned 8% total, 6% from ClimaCore, and 2% from Atmos). If that fails to recover the last few percent, and people are still unhappy with the regression at helem=16, then we can just revert all of ClimaCore back to that commit and reapply changes since.

Zhaoyi is already factoring out the contribution from ClimaAtmos here.

This job has SYPD of 0.892 at week 13. If you run an identical job but with ClimaCore 0.14.12, the SYPD ws 0.956 (7% faster). If you run an identical job but with ClimaCore 0.14.13, the SYPD was 0.869 (2.5% slower).

I'd like to iterate on this more quickly, how can I run this job interactively?

julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file config/amip_configs/amip.yml --job_id amip

[charleskawczynski]

local notes:

# julia --project=experiments/ClimaEarth/
empty!(ARGS)
ENV["CLIMACOMMS_DEVICE"] = "CUDA";
push!(ARGS, "--config_file", "config/amip_configs/amip.yml")
push!(ARGS, "--job_id", "amip")
using Revise; include("experiments/ClimaEarth/run_amip.jl")

[charleskawczynski]

Superseded by #2055.