Merge pull request #294 from CliMA/ck/customize_hcrop

Add horizontal cropping to perf table

Author: charleskawczynski

Project.toml

@@ -1,7 +1,7 @@
 name = "ClimaTimeSteppers"
 uuid = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 authors = ["Climate Modeling Alliance"]
-version = "0.7.30"
+version = "0.7.31"
 
 [deps]
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"

ext/ClimaTimeSteppersBenchmarkToolsExt.jl

@@ -53,6 +53,7 @@ const allowed_names =
         with_cu_prof::Symbol = :bfrofile, # [:bprofile, :profile]
         trace::Bool = false,
         crop::Bool = false,
+        hcrop::Union{Nothing, Int} = nothing,
         only::Union{Nothing, Vector{String}} = nothing,
     )
 
@@ -62,6 +63,7 @@ Benchmark a DistributedODEIntegrator given:
  - `with_cu_prof`, `:profile` or `:bprofile`, to call `CUDA.@profile` or `CUDA.@bprofile` respectively.
  - `trace`, Bool passed to `CUDA.@profile` (see CUDA docs)
  - `crop`, Bool indicating whether or not to crop the `CUDA.@profile` printed table.
+ - `hcrop`, Number of horizontal characters to include in the table before cropping.
  - `only, list of functions to benchmarks (benchmark all by default)
 
 `only` may contain:
@@ -81,6 +83,7 @@ function CTS.benchmark_step(
     with_cu_prof::Symbol = :bprofile,
     trace::Bool = false,
     crop::Bool = false,
+    hcrop::Union{Nothing, Int} = nothing,
     only::Union{Nothing, Vector{String}} = nothing,
 )
     (; u, p, t, dt, sol, alg) = integrator
@@ -96,7 +99,7 @@ function CTS.benchmark_step(
         @. X = u
         @. Xlim = u
         trials₀ = OrderedCollections.OrderedDict()
-        kwargs = (; device, with_cu_prof, trace, crop)
+        kwargs = (; device, with_cu_prof, trace, crop, hcrop)
 #! format: off
         maybe_push!(trials₀, "Wfact", wfact_fun(integrator), (W, u, p, dt, t), kwargs, only)
         maybe_push!(trials₀, "ldiv!", LA.ldiv!, (X, W, u), kwargs, only)

ext/benchmark_utils.jl

@@ -48,8 +48,9 @@ function tabulate_summary(summary; n_calls_per_step)
     )
 end
 
-get_trial(f::Nothing, args, name; device, with_cu_prof = :bprofile, trace = false, crop = false) = nothing
-function get_trial(f, args, name; device, with_cu_prof = :bprofile, trace = false, crop = false)
+get_trial(f::Nothing, args, name; device, with_cu_prof = :bprofile, trace = false, crop = false, hcrop = nothing) =
+    nothing
+function get_trial(f, args, name; device, with_cu_prof = :bprofile, trace = false, crop = false, hcrop = nothing)
     f(args...) # compile first
     b = if device isa ClimaComms.CUDADevice
         BenchmarkTools.@benchmarkable CUDA.@sync $f($(args)...)
@@ -68,8 +69,13 @@ function get_trial(f, args, name; device, with_cu_prof = :bprofile, trace = fals
         if crop
             println(p) # crops by default
         else
-            io = IOContext(stdout, :crop => :horizontal)
-            show(io, p)
+            # use "COLUMNS" to set how many horizontal characters to crop:
+            # See https://github.com/ronisbr/PrettyTables.jl/issues/11#issuecomment-2145550354
+            envs = isnothing(hcrop) ? () : ("COLUMNS" => hcrop,)
+            withenv(envs...) do
+                io = IOContext(stdout, :crop => :horizontal, :limit => true, :displaysize => displaysize())
+                show(io, p)
+            end
             println()
         end
     end