Use local memory in band matrix solve

[charleskawczynski]

This PR adds the use of local memory in the band matrix solve.

Closes #1738.

[charleskawczynski]

Main branch:

┌ Warning: Percentages are only based on ["ldiv!"], pass `only = nothing` for accurately reported percentages
└ @ ClimaTimeSteppersBenchmarkToolsExt ~/CliMA/ClimaTimeSteppers.jl/ext/ClimaTimeSteppersBenchmarkToolsExt.jl:128
[ Info: (#)x entries have been multiplied by corresponding factors in order to compute percentages
┌────────────┬────────────┬──────────┬────────────┬────────────┬────────────┬────────────┬───────────┬──────────────────┐
│ Function   │     Memory │   allocs │       Time │       Time │       Time │       Time │ N-samples │ step! percentage │
│            │   estimate │ estimate │        min │        max │       mean │     median │           │                  │
├────────────┼────────────┼──────────┼────────────┼────────────┼────────────┼────────────┼───────────┼──────────────────┤
│ ldiv! (9x) │ 696.80 KiB │    15363 │ 896.633 ms │ 901.964 ms │ 898.880 ms │ 899.101 ms │        10 │               -1 │
└────────────┴────────────┴──────────┴────────────┴────────────┴────────────┴────────────┴───────────┴──────────────────┘

This branch:

┌ Warning: Percentages are only based on ["ldiv!"], pass `only = nothing` for accurately reported percentages
└ @ ClimaTimeSteppersBenchmarkToolsExt ~/CliMA/ClimaTimeSteppers.jl/ext/ClimaTimeSteppersBenchmarkToolsExt.jl:128
[ Info: (#)x entries have been multiplied by corresponding factors in order to compute percentages
┌────────────┬────────────┬──────────┬────────────┬────────────┬────────────┬────────────┬───────────┬──────────────────┐
│ Function   │     Memory │   allocs │       Time │       Time │       Time │       Time │ N-samples │ step! percentage │
│            │   estimate │ estimate │        min │        max │       mean │     median │           │                  │
├────────────┼────────────┼──────────┼────────────┼────────────┼────────────┼────────────┼───────────┼──────────────────┤
│ ldiv! (9x) │ 816.89 KiB │    17001 │ 377.620 ms │ 453.659 ms │ 395.300 ms │ 377.983 ms │        10 │               -1 │
└────────────┴────────────┴──────────┴────────────┴────────────┴────────────┴────────────┴───────────┴──────────────────┘

So, a 2.37x speedup in ldiv!.

Based on the device-side activity (on this branch):

Device-side activity: GPU was busy for 42.09 ms (98.75% of the trace)
┌──────────┬────────────┬───────┬──────────────────────────────────────┬──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
│ Time (%) │ Total time │ Calls │ Time distribution                    │ Name                                                                                                                                                                                                        ⋯
├──────────┼────────────┼───────┼──────────────────────────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
│   29.91% │   12.75 ms │    12 │   1.06 ms ± 0.0    (  1.06 ‥ 1.07)   │ _Z26single_field_solve_kernel_10CUDADevice5FieldI5VIJFHI5TupleI7Float32S3_ELi4E13CuDeviceArrayIS3_Li5ELi1EEE29ExtrudedFiniteDifferenceSpaceI34DeviceExtrudedFiniteDifferenceGridI22DeviceIntervalTopologyI1 ⋯
│   13.33% │    5.68 ms │    21 │ 270.48 µs ± 0.76   ( 268.7 ‥ 272.04) │ _Z22copyto_stencil_kernel_5FieldI5VIJFHI7Float32Li4E13CuDeviceArrayIS1_Li5ELi1EEE16PlaceholderSpaceE11BroadcastedI22CUDAColumnStencilStyleS3_8identity5TupleIS4_IS5_S3_4rsubS7_IS_IS0_IS1_Li4ES2_IS1_Li5ELi ⋯
│    6.67% │    2.84 ms │     2 │   1.42 ms ± 0.0    (  1.42 ‥ 1.42)   │ _Z26single_field_solve_kernel_10CUDADevice5FieldI5VIJFHI5TupleI10AxisTensorI7Float32Li1ES2_I13CovariantAxisI6_1__2_EE6SArrayIS2_ILi2EES4_Li1ELi2EEES4_ELi4E13CuDeviceArrayIS4_Li5ELi1EEE29ExtrudedFiniteDif ⋯
│    6.21% │    2.65 ms │    11 │  240.8 µs ± 0.32   (240.09 ‥ 241.28) │ _Z22copyto_stencil_kernel_5FieldI5VIJFHI10AxisTensorI7Float32Li1E5TupleI13CovariantAxisI4_3__EE6SArrayIS3_ILi1EES2_Li1ELi1EEELi4E13CuDeviceArrayIS2_Li5ELi1EEE16PlaceholderSpaceE11BroadcastedI22CUDAColumn ⋯
│    6.19% │    2.64 ms │     2 │   1.32 ms ± 0.0    (  1.32 ‥ 1.32)   │ _Z26single_field_solve_kernel_10CUDADevice5FieldI5VIJFHI5TupleI10AxisTensorI7Float32Li1ES2_I13CovariantAxisI4_3__EE6SArrayIS2_ILi1EES4_Li1ELi1EEES3_IS4_Li2ES2_IS5_I4_3__E17ContravariantAxisI4_3__EES6_IS2 ⋯
│    6.01% │    2.56 ms │    42 │   61.0 µs ± 3.09   ( 57.22 ‥ 65.8)   │ _Z11knl_copyto_5VIJFHI7Float32Li4E13CuDeviceArrayIS0_Li5ELi1EEE11BroadcastedI10VIJFHStyleILi4E50CuArray_Float32__N__CUDA_Mem_DeviceBuffer__where_NE5TupleI5OneToI5Int64ES5_IS6_ES5_IS6_ES5_IS6_ES5_IS6_EE8i ⋯
│    3.88% │    1.66 ms │     2 │ 827.67 µs ± 0.17   (827.55 ‥ 827.79) │ _Z22copyto_stencil_kernel_5FieldI5VIJFHI10AxisTensorI7Float32Li1E5TupleI13CovariantAxisI4_3__EE6SArrayIS3_ILi1EES2_Li1ELi1EEELi4E13CuDeviceArrayIS2_Li5ELi1EEE16PlaceholderSpaceE11BroadcastedI22CUDAColumn ⋯
│    2.82% │     1.2 ms │     1 │                                      │ _Z22copyto_stencil_kernel_5FieldI5VIJFHI13BandMatrixRowILin1ELi3E10AxisTensorI7Float32Li2E5TupleI13CovariantAxisI4_3__E17ContravariantAxisI4_3__EE6SArrayIS4_ILi1ELi1EES3_Li2ELi1EEEELi4E13CuDeviceArrayIS3 ⋯
│    2.32% │  990.63 µs │     4 │ 247.66 µs ± 1.0    (246.29 ‥ 248.67) │ _Z28multiple_field_solve_kernel_10CUDADevice5TupleI5FieldI5VIJFHIS0_I7Float32S3_ELi4E13CuDeviceArrayIS3_Li5ELi1EEE16PlaceholderSpaceES1_IS2_IS0_I10AxisTensorIS3_Li1ES0_I13CovariantAxisI4_3__EE6SArrayIS0_ ⋯
│    2.31% │  986.58 µs │     4 │ 246.64 µs ± 0.92   (245.57 ‥ 247.72) │ _Z28multiple_field_solve_kernel_10CUDADevice5TupleI5FieldI5VIJFHIS0_I7Float32S3_ELi4E13CuDeviceArrayIS3_Li5ELi1EEE16PlaceholderSpaceES1_IS2_IS0_IS3_S3_ELi4ES4_IS3_Li5ELi1EEES5_EES0_IS1_IS2_IS3_Li4E8SubAr ⋯
│    2.22% │  947.48 µs │     7 │ 135.35 µs ± 0.27   (135.18 ‥ 135.9)  │ _Z22copyto_stencil_kernel_5FieldI5VIJFHI7Float32Li4E13CuDeviceArrayIS1_Li5ELi1EEE16PlaceholderSpaceE11BroadcastedI22CUDAColumnStencilStyleS3_8identity5TupleIS4_IS5_S3_4rsubS7_IS_IS0_IS1_Li4ES2_IS1_Li5ELi ⋯
│    2.10% │  897.17 µs │     6 │ 149.53 µs ± 0.32   (149.01 ‥ 149.97) │ _Z22copyto_stencil_kernel_5FieldI5VIJFHI7Float32Li4E13CuDeviceArrayIS1_Li5ELi1EEE16PlaceholderSpaceE11BroadcastedI22CUDAColumnStencilStyleS3_8identity5TupleI18StencilBroadcastedIS5_33MultiplyColumnwiseBa ⋯
│    1.94% │  827.79 µs │     3 │ 275.93 µs ± 0.5    (275.37 ‥ 276.33) │ _Z22copyto_stencil_kernel_5FieldI5VIJFHI7Float32Li4E13CuDeviceArrayIS1_Li5ELi1EEE16PlaceholderSpaceE11BroadcastedI22CUDAColumnStencilStyleS3_8identity5TupleIS4_IS5_S3_4rsubS7_IS_IS0_IS1_Li4E8SubArrayIS1_ ⋯
│    1.61% │  685.21 µs │     2 │ 342.61 µs ± 0.67   (342.13 ‥ 343.08) │ _Z22copyto_stencil_kernel_5FieldI5VIJFHI10AxisTensorI7Float32Li1E5TupleI13CovariantAxisI6_1__2_EE6SArrayIS3_ILi2EES2_Li1ELi2EEELi4E13CuDeviceArrayIS2_Li5ELi1EEE16PlaceholderSpaceE11BroadcastedI22CUDAColu ⋯
│    1.36% │  579.36 µs │     2 │ 289.68 µs ± 0.0    (289.68 ‥ 289.68) │ _Z22copyto_stencil_kernel_5FieldI5VIJFHI10AxisTensorI7Float32Li1E5TupleI13CovariantAxisI4_3__EE6SArrayIS3_ILi1EES2_Li1ELi1EEELi4E13CuDeviceArrayIS2_Li5ELi1EEE16PlaceholderSpaceE11BroadcastedI22CUDAColumn ⋯
│    1.12% │  476.12 µs │     1 │                                      │ _Z22copyto_stencil_kernel_5FieldI5VIJFHI10AxisTensorI7Float32Li1E5TupleI13CovariantAxisI4_3__EE6SArrayIS3_ILi1EES2_Li1ELi1EEELi4E13CuDeviceArrayIS2_Li5ELi1EEE16PlaceholderSpaceE11BroadcastedI22CUDAColumn ⋯
│    1.08% │  458.96 µs │     3 │ 152.99 µs ± 0.14   (152.83 ‥ 153.06) │ _Z22copyto_stencil_kernel_5FieldI5VIJFHI7Float32Li4E13CuDeviceArrayIS1_Li5ELi1EEE16PlaceholderSpaceE11BroadcastedI22CUDAColumnStencilStyleS3_8identity5TupleIS4_IS5_S3_4rsubS7_IS_IS0_IS1_Li4E8SubArrayIS1_ ⋯
│    1.01% │  430.35 µs │     2 │ 215.17 µs ± 0.51   (214.82 ‥ 215.53) │ _Z11knl_copyto_5VIJFHI10AxisTensorI7Float32Li1E5TupleI13CovariantAxisI6_1__2_EE6SArrayIS2_ILi2EES1_Li1ELi2EEELi4E8SubArrayIS1_Li5E13CuDeviceArrayIS1_Li5ELi1EES2_I5SliceI5OneToI5Int64EES7_IS8_IS9_EES7_IS8 ⋯
│    0.92% │  391.96 µs │     8 │   49.0 µs ± 3.48   ( 42.92 ‥ 51.5)   │ _Z11knl_copyto_5VIJFHI7Float32Li4E13CuDeviceArrayIS0_Li5ELi1EEE11BroadcastedI10VIJFHStyleILi4E50CuArray_Float32__N__CUDA_Mem_DeviceBuffer__where_NE5TupleI5OneToI5Int64ES5_IS6_ES5_IS6_ES5_IS6_ES5_IS6_EE8i ⋯
│    0.91% │  387.43 µs │     6 │  64.57 µs ± 3.53   ( 60.32 ‥ 68.43)  │ _Z11knl_copyto_5VIJFHI7Float32Li4E13CuDeviceArrayIS0_Li5ELi1EEE11BroadcastedI10VIJFHStyleILi4E50CuArray_Float32__N__CUDA_Mem_DeviceBuffer__where_NE5TupleI5OneToI5Int64ES5_IS6_ES5_IS6_ES5_IS6_ES5_IS6_EE8i ⋯
│    0.79% │  338.79 µs │     1 │                                      │ _Z11knl_copyto_5VIJFHI13BandMatrixRowILin1ELi3E7Float32ELi4E13CuDeviceArrayIS1_Li5ELi1EEE11BroadcastedI10VIJFHStyleILi4E50CuArray_Float32__N__CUDA_Mem_DeviceBuffer__where_NE5TupleI5OneToI5Int64ES6_IS7_ES ⋯
│    0.74% │  316.38 µs │     1 │                                      │ _Z11knl_copyto_5VIJFHI13BandMatrixRowI39ClimaCore_Utilities_PlusHalf_Int64___1_Li2E10AxisTensorI7Float32Li2E5TupleI13CovariantAxisI4_3__E17ContravariantAxisI6_1__2_EE6SArrayIS3_ILi1ELi2EES2_Li2ELi2EEEELi ⋯
│    0.59% │  252.96 µs │     3 │  84.32 µs ± 1.93   ( 82.25 ‥ 86.07)  │ _Z11knl_copyto_5VIJFHI7Float32Li4E8SubArrayIS0_Li5E13CuDeviceArrayIS0_Li5ELi1EE5TupleI5SliceI5OneToI5Int64EES4_IS5_IS6_EES4_IS5_IS6_EE9UnitRangeIS6_ES4_IS5_IS6_EEELinfalseEEE11BroadcastedI10VIJFHStyleILi ⋯
│    0.58% │  245.57 µs │     1 │                                      │ _Z22copyto_stencil_kernel_5FieldI5VIJFHI10AxisTensorI7Float32Li1E5TupleI13CovariantAxisI4_3__EE6SArrayIS3_ILi1EES2_Li1ELi1EEELi4E13CuDeviceArrayIS2_Li5ELi1EEE16PlaceholderSpaceE11BroadcastedI22CUDAColumn ⋯
│    0.50% │  212.43 µs │     1 │                                      │ _Z11knl_copyto_5VIJFHI10AxisTensorI7Float32Li1E5TupleI13CovariantAxisI6_1__2_EE6SArrayIS2_ILi2EES1_Li1ELi2EEELi4E8SubArrayIS1_Li5E13CuDeviceArrayIS1_Li5ELi1EES2_I5SliceI5OneToI5Int64EES7_IS8_IS9_EES7_IS8 ⋯
│    0.42% │  176.91 µs │     1 │                                      │ _Z11knl_copyto_5VIJFHI10AxisTensorI7Float32Li1E5TupleI13CovariantAxisI6_1__2_EE6SArrayIS2_ILi2EES1_Li1ELi2EEELi4E13CuDeviceArrayIS1_Li5ELi1EEE11BroadcastedI10VIJFHStyleILi4E50CuArray_Float32__N__CUDA_Mem ⋯
│    0.34% │  143.29 µs │     1 │                                      │ _Z22copyto_stencil_kernel_5FieldI5VIJFHI7Float32Li4E13CuDeviceArrayIS1_Li5ELi1EEE16PlaceholderSpaceE11BroadcastedI22CUDAColumnStencilStyleS3_8identity5TupleIS4_IS5_S3_4rsubS7_IS_IS0_IS1_Li4E8SubArrayIS1_ ⋯
│    0.31% │  132.32 µs │     2 │  66.16 µs ± 0.17   ( 66.04 ‥ 66.28)  │ _Z11knl_copyto_5VIJFHI10AxisTensorI7Float32Li1E5TupleI13CovariantAxisI4_3__EE6SArrayIS2_ILi1EES1_Li1ELi1EEELi4E8SubArrayIS1_Li5E13CuDeviceArrayIS1_Li5ELi1EES2_I5SliceI5OneToI5Int64EES7_IS8_IS9_EES7_IS8_I ⋯
│    0.18% │   78.68 µs │     1 │                                      │ _Z11knl_copyto_5VIJFHI7Float32Li4E8SubArrayIS0_Li5E13CuDeviceArrayIS0_Li5ELi1EE5TupleI5SliceI5OneToI5Int64EES4_IS5_IS6_EES4_IS5_IS6_EE9UnitRangeIS6_ES4_IS5_IS6_EEELinfalseEEE11BroadcastedI10VIJFHStyleILi ⋯
│    0.16% │   66.52 µs │     2 │  33.26 µs ± 0.17   ( 33.14 ‥ 33.38)  │ _Z11knl_copyto_5VIJFHI10AxisTensorI7Float32Li1E5TupleI13CovariantAxisI4_3__EE6SArrayIS2_ILi1EES1_Li1ELi1EEELi4E13CuDeviceArrayIS1_Li5ELi1EEE11BroadcastedI10VIJFHStyleILi4E50CuArray_Float32__N__CUDA_Mem_D ⋯
│    0.12% │   49.11 µs │     1 │                                      │ _Z11knl_copyto_5VIJFHI10AxisTensorI7Float32Li1E5TupleI13CovariantAxisI4_3__EE6SArrayIS2_ILi1EES1_Li1ELi1EEELi4E13CuDeviceArrayIS1_Li5ELi1EEE11BroadcastedI10VIJFHStyleILi4E50CuArray_Float32__N__CUDA_Mem_D ⋯
│    0.08% │   35.52 µs │     1 │                                      │ _Z11knl_copyto_5VIJFHI10AxisTensorI7Float32Li1E5TupleI13CovariantAxisI4_3__EE6SArrayIS2_ILi1EES1_Li1ELi1EEELi4E8SubArrayIS1_Li5E13CuDeviceArrayIS1_Li5ELi1EES2_I5SliceI5OneToI5Int64EES7_IS8_IS9_EES7_IS8_I ⋯
└──────────┴────────────┴───────┴──────────────────────────────────────┴──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
                                                                                                                                                                                                                                                                      1 column omitted

It seems like the next step is to use a better (parallel) algorithm for the tridiagonal band matrix solve, to see if we can get better than ~1ms.