Merge pull request #1836 from CliMA/ck/misc

Update and add cuda stencil benchmarks

Author: charleskawczynski

.buildkite/pipeline.yml

@@ -1180,15 +1180,15 @@ steps:
         key: "perf_cpu_implicit_stencil"
         command: "julia --color=yes --check-bounds=yes --project=.buildkite test/Operators/finitedifference/opt_implicit_stencils.jl"
 
-      - label: "Perf: FD operator benchmarks"
+      - label: "Perf: FD operator stencil benchmarks"
         key: "perf_fd_ops"
-        command: "julia --color=yes --project=.buildkite test/Operators/finitedifference/benchmark_column.jl"
+        command: "julia --color=yes --project=.buildkite test/Operators/finitedifference/benchmark_stencils.jl"
 
-      - label: "Perf: GPU FD operator benchmarks"
+      - label: "Perf: GPU FD operator stencil benchmarks"
         key: "gpu_perf_fd_ops"
         command:
           - "julia --project=.buildkite -e 'using CUDA; CUDA.versioninfo()'"
-          - "julia --color=yes --project=.buildkite test/Operators/finitedifference/benchmark_column.jl"
+          - "julia --color=yes --project=.buildkite test/Operators/finitedifference/benchmark_stencils.jl"
         env:
           CLIMACOMMS_DEVICE: "CUDA"
         agents:

test/Operators/finitedifference/benchmark_column.jl renamed to test/Operators/finitedifference/benchmark_stencils.jl

@@ -1,8 +1,8 @@
 #=
 julia --project
-using Revise; include(joinpath("test", "Operators", "finitedifference", "benchmark_column.jl"))
+using Revise; include(joinpath("test", "Operators", "finitedifference", "benchmark_stencils.jl"))
 =#
-include("benchmark_column_utils.jl")
+include("benchmark_stencils_utils.jl")
 
 @testset "Benchmark operators" begin
     benchmark_operators(Float64; z_elems = 63, helem = 30, Nq = 4)

test/Operators/finitedifference/benchmark_stencils_array_kernels.jl

@@ -0,0 +1,131 @@
+import CUDA
+#####
+##### CPU column
+#####
+
+#= e.g., any 2nd order interpolation / derivative operator =#
+function column_op_2mul_1add!(x, y, D, U)
+    y1 = @view y[1:(end - 1)]
+    y2 = @view y[2:end]
+    @inbounds for i in eachindex(x)
+        x[i] = D[i] * y1[i] + U[i] * y2[i]
+    end
+    return nothing
+end
+
+#= e.g., div(grad(scalar)), div(interp(vec)) =#
+function column_op_3mul_2add!(x, y, L, D, U)
+    y1 = @view y[1:(end - 1)]
+    y2 = @view y[2:(end - 1)]
+    y3 = @view y[2:end]
+    @inbounds for i in eachindex(x)
+        i == 1 && continue
+        i == length(x) && continue
+        x[i] = L[i] * y1[i] + D[i] * y2[i] + U[i] * y3[i]
+    end
+    return nothing
+end
+
+#= e.g., curlC2F =#
+function column_curl_like!(curluₕ, uₕ_x, uₕ_y, D, U)
+    @inbounds for i in eachindex(curluₕ)
+        curluₕ[i] = D[i] * uₕ_x[i] + U[i] * uₕ_y[i]
+    end
+    return nothing
+end
+
+#####
+##### CUDA column
+#####
+
+# TODO: expand on this
+
+function column_op_2mul_1add_cuda!(x, y, D, U)
+    kernel =
+        CUDA.@cuda always_inline = true launch = false op_2mul_1add_cuda_kernel!(
+            x,
+            y,
+            D,
+            U,
+            Val(length(x)),
+        )
+    config = CUDA.launch_configuration(kernel.fun)
+    nitems = length(x)
+    threads = min(nitems, config.threads)
+    blocks = cld(nitems, threads)
+    kernel(x, y, D, U; threads, blocks) # This knows to use always_inline from above.
+    return nothing
+end
+
+function op_2mul_1add_cuda_kernel!(x, y, D, U, ::Val{N}) where {N}
+    @inbounds begin
+        i = thread_index()
+        if i ≤ N
+            x[i] = D[i] * y[i] + U[i] * y[i + 1]
+        end
+    end
+    return nothing
+end
+
+
+#####
+##### CPU sphere
+#####
+
+# TODO
+
+#####
+##### CUDA sphere
+#####
+
+# TODO: move to CUDA utils
+thread_index() =
+    (CUDA.blockIdx().x - Int32(1)) * CUDA.blockDim().x + CUDA.threadIdx().x
+Base.@propagate_inbounds kernel_indexes(tidx, Nv, Nij, Nh) =
+    CartesianIndices(map(x -> Base.OneTo(x), (Nv, Nij, Nij, 1, Nh)))[tidx]
+valid_range(tidx, n) = 1 ≤ tidx ≤ n
+
+
+#= e.g., any 2nd order interpolation / derivative operator =#
+function sphere_op_2mul_1add_cuda!(x, y, D, U)
+    Nv = size(x, 1)
+    Nij = size(x, 2)
+    Nh = size(x, 5)
+    N = length(x)
+    kernel =
+        CUDA.@cuda always_inline = true launch = false sphere_op_2mul_1add_cuda_kernel!(
+            x,
+            y,
+            D,
+            U,
+            Val(Nv),
+            Val(Nij),
+            Val(Nh),
+            Val(N),
+        )
+    config = CUDA.launch_configuration(kernel.fun)
+    threads = min(N, config.threads)
+    blocks = cld(N, threads)
+    kernel(x, y, D, U; threads, blocks)
+    return nothing
+end
+
+function sphere_op_2mul_1add_cuda_kernel!(
+    x,
+    y,
+    D,
+    U,
+    ::Val{Nv},
+    ::Val{Nij},
+    ::Val{Nh},
+    ::Val{N},
+) where {Nv, Nij, Nh, N}
+    @inbounds begin
+        tidx = thread_index()
+        if valid_range(tidx, N)
+            I = kernel_indexes(tidx, Nv, Nij, Nh)
+            x[I] = D[I] * y[I] + U[I] * y[I + CartesianIndex(1, 0, 0, 0, 0)]
+        end
+    end
+    return nothing
+end

test/Operators/finitedifference/benchmark_column_kernels.jl renamed to test/Operators/finitedifference/benchmark_stencils_climacore_kernels.jl

@@ -33,41 +33,6 @@ field_vars(::Type{FT}) where {FT} = (;
     ᶠw = Geometry.Covariant3Vector(FT(0)),
 )
 
-#####
-##### Second order interpolation / derivatives
-#####
-
-#= e.g., any 2nd order interpolation / derivative operator =#
-function op_2mul_1add!(x, y, D, U)
-    y1 = @view y[1:(end - 1)]
-    y2 = @view y[2:end]
-    @inbounds for i in eachindex(x)
-        x[i] = D[i] * y1[i] + U[i] * y2[i]
-    end
-    return nothing
-end
-
-#= e.g., div(grad(scalar)), div(interp(vec)) =#
-function op_3mul_2add!(x, y, L, D, U)
-    y1 = @view y[1:(end - 1)]
-    y2 = @view y[2:(end - 1)]
-    y3 = @view y[2:end]
-    @inbounds for i in eachindex(x)
-        i==1 && continue
-        i==length(x) && continue
-        x[i] = L[i] * y1[i] + D[i] * y2[i] + U[i] * y3[i]
-    end
-    return nothing
-end
-
-#= e.g., curlC2F =#
-function curl_like!(curluₕ, uₕ_x, uₕ_y, D, U)
-    @inbounds for i in eachindex(curluₕ)
-        curluₕ[i] = D[i] * uₕ_x[i] + U[i] * uₕ_y[i]
-    end
-    return nothing
-end
-
 function set_value_bcs(c)
     FT = Spaces.undertype(axes(c))
     return (;bottom = Operators.SetValue(FT(0)),
@@ -211,7 +176,8 @@ bc_name(bcs::Tuple) = (:none,)
 bc_name_base(bcs::@NamedTuple{}) = (:none,)
 bc_name(bcs::@NamedTuple{}) = (:none,)
 
-include("benchmark_column_kernels.jl")
+include("benchmark_stencils_array_kernels.jl")
+include("benchmark_stencils_climacore_kernels.jl")
 
 uses_bycolumn(::typeof(op_broadcast_example0!)) = true
 uses_bycolumn(::typeof(op_broadcast_example1!)) = true
@@ -268,54 +234,88 @@ function benchmark_func!(t_min, trials, fun, c, f, verbose = false)
     end
 end
 
-function benchmark_arrays(z_elems, ::Type{FT}) where {FT}
-    L = zeros(FT, z_elems)
-    D = zeros(FT, z_elems)
-    U = zeros(FT, z_elems)
-    xarr = rand(FT, z_elems)
-    uₕ_x = rand(FT, z_elems)
-    uₕ_y = rand(FT, z_elems)
-    yarr = rand(FT, z_elems + 1)
-
-    println("\n############################ 2-point stencil")
-    trial = BenchmarkTools.@benchmark op_2mul_1add!($xarr, $yarr, $D, $U)
-    show(stdout, MIME("text/plain"), trial)
-    println()
-    println("\n############################ 3-point stencil")
-    trial = BenchmarkTools.@benchmark op_3mul_2add!($xarr, $yarr, $L, $D, $U)
-    show(stdout, MIME("text/plain"), trial)
-    println()
-    println("\n############################ curl-like stencil")
-    trial = BenchmarkTools.@benchmark curl_like!($xarr, $uₕ_x, $uₕ_y, $D, $U)
-    show(stdout, MIME("text/plain"), trial)
-    println()
+function column_benchmark_arrays(device, z_elems, ::Type{FT}) where {FT}
+    ArrayType = ClimaComms.array_type(device)
+    L = ArrayType(zeros(FT, z_elems))
+    D = ArrayType(zeros(FT, z_elems))
+    U = ArrayType(zeros(FT, z_elems))
+    xarr = ArrayType(rand(FT, z_elems))
+    uₕ_x = ArrayType(rand(FT, z_elems))
+    uₕ_y = ArrayType(rand(FT, z_elems))
+    yarr = ArrayType(rand(FT, z_elems + 1))
+
+    if device isa ClimaComms.CUDADevice
+        println("\n############################ column 2-point stencil")
+        trial = BenchmarkTools.@benchmark ClimaComms.@cuda_sync $device column_op_2mul_1add_cuda!($xarr, $yarr, $D, $U)
+        show(stdout, MIME("text/plain"), trial)
+        println()
+    else
+        println("\n############################ column 2-point stencil")
+        trial = BenchmarkTools.@benchmark column_op_2mul_1add!($xarr, $yarr, $D, $U)
+        show(stdout, MIME("text/plain"), trial)
+        println()
+        println("\n############################ column 3-point stencil")
+        trial = BenchmarkTools.@benchmark column_op_3mul_2add!($xarr, $yarr, $L, $D, $U)
+        show(stdout, MIME("text/plain"), trial)
+        println()
+        println("\n############################ column curl-like stencil")
+        trial = BenchmarkTools.@benchmark column_curl_like!($xarr, $uₕ_x, $uₕ_y, $D, $U)
+        show(stdout, MIME("text/plain"), trial)
+        println()
+    end
+end
+
+function sphere_benchmark_arrays(device, z_elems, helem, Nq, ::Type{FT}) where {FT}
+    ArrayType = ClimaComms.array_type(device)
+    # VIJFH
+    Nh = helem * helem * 6
+    cdims = (z_elems  , Nq, Nq, 1, Nh)
+    fdims = (z_elems+1, Nq, Nq, 1, Nh)
+    L = ArrayType(zeros(FT, cdims...))
+    D = ArrayType(zeros(FT, cdims...))
+    U = ArrayType(zeros(FT, cdims...))
+    xarr = ArrayType(rand(FT, cdims...))
+    uₕ_x = ArrayType(rand(FT, cdims...))
+    uₕ_y = ArrayType(rand(FT, cdims...))
+    yarr = ArrayType(rand(FT, fdims...))
+
+    if device isa ClimaComms.CUDADevice
+        println("\n############################ sphere 2-point stencil")
+        trial = BenchmarkTools.@benchmark ClimaComms.@cuda_sync $device sphere_op_2mul_1add_cuda!($xarr, $yarr, $D, $U)
+        show(stdout, MIME("text/plain"), trial)
+        println()
+    else
+        @info "Sphere CPU kernels have not been added yet."
+    end
 end
 
 function benchmark_operators(::Type{FT}; z_elems, helem, Nq) where {FT}
-    @show ClimaComms.device()
+    device = ClimaComms.device()
+    @show device
     trials = OrderedCollections.OrderedDict()
     t_min = OrderedCollections.OrderedDict()
-    benchmark_arrays(z_elems, FT)
+    column_benchmark_arrays(device, z_elems, FT)
+    sphere_benchmark_arrays(device, z_elems, helem, Nq, FT)
 
     cspace = TU.ColumnCenterFiniteDifferenceSpace(FT; zelem=z_elems)
     fspace = Spaces.FaceFiniteDifferenceSpace(cspace)
     cfield = fill(field_vars(FT), cspace)
     ffield = fill(field_vars(FT), fspace)
-    benchmark_operators_base(trials, t_min, cfield, ffield)
+    benchmark_operators_base(trials, t_min, cfield, ffield, "column")
 
     cspace = TU.CenterExtrudedFiniteDifferenceSpace(FT; zelem=z_elems, helem, Nq)
     fspace = Spaces.FaceExtrudedFiniteDifferenceSpace(cspace)
     cfield = fill(field_vars(FT), cspace)
     ffield = fill(field_vars(FT), fspace)
-    benchmark_operators_base(trials, t_min, cfield, ffield)
+    benchmark_operators_base(trials, t_min, cfield, ffield, "sphere")
 
     # Tests are removed since they're flakey. And maintaining
     # them before they're converged is a bit of work..
     test_results(t_min)
     return (; trials, t_min)
 end
 
-function benchmark_operators_base(trials, t_min, cfield, ffield)
+function benchmark_operators_base(trials, t_min, cfield, ffield, name)
     ops = [
         #### Core discrete operators
         op_GradientF2C!,
@@ -346,7 +346,7 @@ function benchmark_operators_base(trials, t_min, cfield, ffield)
         op_divgrad_uₕ!,
     ]
 
-    @info "Benchmarking operators, this may take a minute or two..."
+    @info "Benchmarking $name operators, this may take a minute or two..."
     for op in ops
         if uses_bycolumn(op) && axes(cfield) isa Spaces.FiniteDifferenceSpace
             continue

test/Operators/finitedifference/benchmark_column_utils.jl renamed to test/Operators/finitedifference/benchmark_stencils_utils.jl

@@ -1,5 +1,5 @@
 #= This script is helpful for targeting specific kernels =#
-include("benchmark_column_utils.jl")
+include("benchmark_stencils_utils.jl")
 
 function apply_kernel!(cfield, ffield)
     # op_DivergenceF2C!(cfield, ffield)