Merge #538

538: Update perf/allocations analysis tools r=charleskawczynski a=charleskawczynski



Co-authored-by: Charles Kawczynski <[email protected]>

.buildkite/pipeline.yml

@@ -156,6 +156,17 @@ steps:
       queue: central
       slurm_ntasks: 1
 
+  - label: ":rocket: Allocations analysis"
+    key: "cpu_allocations"
+    command:
+      - "julia --color=yes --project=test --track-allocation=user perf/allocs.jl"
+    artifact_paths:
+      - "perf/allocations_output/*"
+    agents:
+      config: cpu
+      queue: central
+      slurm_ntasks: 1
+
   - wait: ~
     continue_on_failure: true
 

perf/alloc_per_case.jl

@@ -0,0 +1,19 @@
+# Launch with `julia --project --track-allocation=user`
+if !("." in LOAD_PATH)
+    push!(LOAD_PATH, ".")
+end
+import Profile
+
+include("common.jl")
+case_name = ENV["ALLOCATION_CASE_NAME"]
+@info "Recording allocations for $case_name"
+sim = init_sim(case_name)
+update_n(sim, 1) # compile first
+Profile.clear_malloc_data()
+update_n(sim, 1)
+
+# Quit julia (which generates .mem files), then call
+#=
+import Coverage
+allocs = Coverage.analyze_malloc("src")
+=#

perf/allocs.jl

@@ -0,0 +1,97 @@
+if !("." in LOAD_PATH)
+    push!(LOAD_PATH, ".")
+end
+import Coverage
+import Plots
+
+# https://github.com/jheinen/GR.jl/issues/278#issuecomment-587090846
+ENV["GKSwstype"] = "nul"
+
+all_cases = [
+    # "ARM_SGP",
+    # "Bomex",
+    # "DryBubble",
+    # "DYCOMS_RF01",
+    # "GABLS",
+    # "GATE_III",
+    # "life_cycle_Tan2018",
+    # "Nieuwstadt",
+    "Rico",
+    # "Soares",
+    # "SP",
+    "TRMM_LBA",
+    "LES_driven_SCM",
+]
+
+filter!(x -> x ≠ "GATE_III", all_cases) # no mse tables for GATE_III
+filter!(x -> x ≠ "SP", all_cases) # not currently running SP
+allocs = Dict()
+for case in all_cases
+    ENV["ALLOCATION_CASE_NAME"] = case
+    run(`julia --project=test/ --track-allocation=user perf/alloc_per_case.jl`)
+    allocs[case] = Coverage.analyze_malloc(".")
+
+    # Clean up files
+    all_files = [joinpath(root, f) for (root, dirs, files) in Base.Filesystem.walkdir(".") for f in files]
+    all_mem_files = filter(x -> endswith(x, ".mem"), all_files)
+    for f in all_mem_files
+        rm(f)
+    end
+end
+
+@info "Post-processing allocations"
+
+function plot_allocs(case_name, allocs_per_case, n_unique_bytes)
+    p = Plots.plot()
+    case_bytes = getproperty.(allocs_per_case, :bytes)[end:-1:1]
+    case_filename = getproperty.(allocs_per_case, :filename)[end:-1:1]
+    case_linenumber = getproperty.(allocs_per_case, :linenumber)[end:-1:1]
+    all_bytes = Int[]
+    filenames = String[]
+    linenumbers = Int[]
+    loc_ids = String[]
+    for (bytes, filename, linenumber) in zip(case_bytes, case_filename, case_linenumber)
+        filename_only = first(split(filename, ".jl")) * ".jl"
+        if endswith(filename_only, "TurbulenceConvection.jl") && linenumber == 1
+            continue # Skip loading module
+        end
+        loc_id = "$filename_only" * "$linenumber"
+        if !(bytes in all_bytes) && !(loc_id in loc_ids)
+            push!(all_bytes, bytes)
+            push!(filenames, filename)
+            push!(linenumbers, linenumber)
+            push!(loc_ids, loc_id)
+            if length(all_bytes) ≥ n_unique_bytes
+                break
+            end
+        end
+    end
+
+    all_bytes = all_bytes ./ 10^3
+    max_bytes = maximum(all_bytes)
+    @info "$case_name: $all_bytes"
+    xtick_name(filename, linenumber) = "$filename, line number: $linenumber"
+    markershape = (:circle, :star, :square, :hexagon)
+    for (bytes, filename, linenumber) in zip(all_bytes, filenames, linenumbers)
+        markershape = (markershape[end], markershape[1:(end - 1)]...)
+        filename_only = first(split(filename, ".jl")) * ".jl"
+        Plots.plot!(
+            [0],
+            [bytes];
+            seriestype = :scatter,
+            label = xtick_name(filename_only, linenumber),
+            markershape = markershape[1],
+            markersize = 1 + bytes / max_bytes * 10,
+        )
+    end
+    Plots.plot!(ylabel = "Number of allocations (KB)")
+    Plots.savefig(joinpath(folder, "allocations_$case_name.png"))
+end
+
+folder = "perf/allocations_output"
+mkpath(folder)
+
+@info "Allocated bytes for single tendency per case:"
+for case in all_cases
+    plot_allocs(case, allocs[case], 10)
+end

perf/common.jl

@@ -0,0 +1,41 @@
+import TurbulenceConvection
+
+tc_dir = dirname(dirname(pathof(TurbulenceConvection)))
+include(joinpath(tc_dir, "integration_tests", "utils", "generate_namelist.jl"))
+include(joinpath(tc_dir, "integration_tests", "utils", "Cases.jl"))
+include(joinpath(tc_dir, "integration_tests", "utils", "parameter_set.jl"))
+include(joinpath(tc_dir, "integration_tests", "utils", "main.jl"))
+import .NameList
+
+TurbulenceConvection.initialize_io(sim::Simulation1d) = nothing
+TurbulenceConvection.io(sim::Simulation1d) = nothing
+
+update_n(sim, N::Int) = update_n(sim, Val(N))
+
+function update_n(sim, ::Val{N}) where {N}
+    for i in 1:N
+        TC.update(sim.Turb, sim.grid, sim.state, sim.GMV, sim.Case, sim.TS)
+    end
+    return nothing
+end
+
+function init_sim(case_name)
+    @info "Initializing $case_name for single timestep, with no IO."
+    @info "call update_n(sim, n) to run update n-times"
+    namelist = NameList.default_namelist(case_name)
+    namelist["time_stepping"]["t_max"] = namelist["time_stepping"]["dt"]
+    namelist["stats_io"]["frequency"] = 10.0e10
+    namelist["stats_io"]["skip"] = true
+    namelist["meta"]["uuid"] = "01"
+    sim = Simulation1d(namelist)
+
+    Cases.initialize_profiles(sim.Case, sim.grid, sim.GMV, sim.state)
+    TC.satadjust(sim.GMV, sim.grid, sim.state)
+
+    Cases.initialize_surface(sim.Case, sim.grid, sim.state, sim.param_set)
+    Cases.initialize_forcing(sim.Case, sim.grid, sim.state, sim.GMV, sim.param_set)
+    Cases.initialize_radiation(sim.Case, sim.grid, sim.state, sim.GMV, sim.param_set)
+
+    initialize_edmf(sim.Turb, sim.grid, sim.state, sim.Case, sim.GMV, sim.TS)
+    return sim
+end

perf/perf.jl

@@ -0,0 +1,6 @@
+include("common.jl")
+import BenchmarkTools
+
+sim = init_sim("Bomex")
+update_n(sim, 1) # compile first
+BenchmarkTools.@benchmark update_n($sim, 1)

perf/profile.jl

@@ -0,0 +1,10 @@
+include("common.jl")
+import ProfileView
+import Profile
+
+Profile.@profile update_n(sim, 100)
+Profile.print()
+# Profile.print(; format = :flat, sortedby = :count)
+ProfileView.@profview update_n(sim, 1000) # compile first
+ProfileView.@profview update_n(sim, 1000)
+Profile.clear_malloc_data()

test/Project.toml

@@ -1,9 +1,11 @@
 [deps]
 ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
 ArtifactWrappers = "a14bc488-3040-4b00-9dc1-f6467924858a"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CLIMAParameters = "6eacf6c3-8458-43b9-ae03-caf5306d3d53"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
 CloudMicrophysics = "6a9e3e04-43cd-43ba-94b9-e8782df3c71b"
+Coverage = "a2441757-f6aa-5fb2-8edb-039e3f45d037"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Debugger = "31a5f54b-26ea-5ae9-a837-f05ce5417438"
 Dierckx = "39dd38d3-220a-591b-8e3c-4c3a8c710a94"
@@ -23,6 +25,7 @@ Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 PoissonRandom = "e409e4f3-bfea-5376-8464-e040bb5c01ab"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RootSolvers = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"