Improvements in benchmark of agentic cuda example

Author: emmanuellujan

examples/agentic/generate-cuda-linear-solver/Project.toml

@@ -1,6 +1,8 @@
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 SmartSolve = "4fbb3a3c-2fa1-4c19-8d57-bae8bc1e16ac"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 

examples/agentic/generate-cuda-linear-solver/benchmark.jl

@@ -2,56 +2,133 @@ using LinearAlgebra
 using SparseArrays
 using CUDA
 using BenchmarkTools
+using OrderedCollections
+using Plots
 
-println("Simple GPU benchmark:\n")
+println("GPU benchmark with error-vs-time plot:\n")
 
 include("solver.jl")
 
-# Problem setup on CPU (sparse)
+# Configuration
 N = 10_000
-A = sprand(N, N, 0.3)
-b = rand(N)
-Ad = CuArray(Matrix(A))
-bd = CuArray(b)
-x2_d = CuArray(zeros(N))
-
-# --- warm-up (compile kernels, allocate, etc.) ---
-x1_d = Ad \ bd
-CUDA.CUSOLVER.gesv!(x2_d, Ad, bd)
-CUDA.synchronize()
-bd = CuArray(b) # reset rhs
-x3_d = proposed_fn(Ad, bd)
-CUDA.synchronize()
-
-println("GPU benchmark (default solver):")
-display(@benchmark begin
-    x = $Ad \ $bd
-    CUDA.synchronize()
-end seconds = 10)
-
-println("\nGPU benchmark (gesv! solver):")
-display(@benchmark begin
-    CUDA.CUSOLVER.gesv!(x2_d, Ad, bd)
-    CUDA.synchronize()
-end seconds = 10)
-
-bd = CuArray(b)
-println("\nGPU benchmark (generated solver):")
-display(@benchmark begin
-    x = proposed_fn($Ad, $bd)
-    CUDA.synchronize()
-end seconds = 10)
-
-# --- error computation on GPU ---
-x1_d = Ad \ bd
-CUDA.CUSOLVER.gesv!(x2_d, Ad, bd)
-x3_d = proposed_fn(Ad, bd)
-CUDA.synchronize()
-
-e1 = norm(Ad * x1_d - bd)
-e2 = norm(Ad * x2_d - bd)
-e3 = norm(Ad * x3_d - bd)
-
-println("\nError of default GPU solver: $e1")
-println("Error of gesv! GPU solver: $e2")
-println("Error of generated GPU solver: $e3")
+sparsity_levels = [0.1, 0.5, 0.9]
+solvers = OrderedDict(
+    "Default" => (Ad, bd) -> (Ad \ bd),
+    "gesv!" => (Ad, bd) -> begin
+        x = CuArray(zeros(size(Ad, 1)))
+        CUDA.CUSOLVER.gesv!(x, Ad, bd)
+        x
+    end,
+    "Generated" => (Ad, bd) -> proposed_fn(Ad, bd)
+)
+
+# Store results for plotting
+results = Dict()
+
+for sparsity in sparsity_levels
+    println("\n=== Sparsity: $sparsity ===")
+    
+    # Generate problem
+    A = sprand(N, N, sparsity)
+    b = rand(N)
+    Ad = CuArray(Matrix(A))
+    bd = CuArray(b)
+    
+    results[sparsity] = Dict()
+    
+    for (solver_name, solver_fn) in solvers
+        println("  $solver_name...")
+        
+        # Warm-up
+        bd_warm = CuArray(copy(b))
+        try
+            x_warm = solver_fn(Ad, bd_warm)
+            CUDA.synchronize()
+        catch e
+            println("    Warning: solver failed during warm-up: $e")
+            continue
+        end
+        
+        # Benchmark
+        bd_bench = CuArray(copy(b))
+        try
+            bench = @benchmark begin
+                x = $(solver_fn)($Ad, $bd_bench)
+                CUDA.synchronize()
+            end seconds = 5 samples = 10
+            
+            time_ms = median(bench.times) / 1e6  # Convert to ms
+            
+            # Compute error
+            bd_err = CuArray(copy(b))
+            x_sol = solver_fn(Ad, bd_err)
+            CUDA.synchronize()
+            error = norm(Ad*x_sol - bd_err) / norm(bd_err)
+            
+            results[sparsity][solver_name] = (time=time_ms, error=error)
+            println("    Time: $(round(time_ms, digits=3)) ms, Error: $(round(error, sigdigits=3))")
+        catch e
+            println("    Error during benchmark: $e")
+        end
+    end
+end
+
+# Create error-vs-time plot
+p = plot(
+    size=(800, 800),
+    #legend=:topright,
+    legend=:bottomright,
+    xlabel="Time (ms)",
+    ylabel="Relative residual: ||Ax - b||₂ / ||b||₂",
+ #   xscale=:log10,
+    yscale=:log10,
+    guidefontsize=22,#18,
+    tickfontsize=20, #16,
+    legendfontsize=18, #14,
+    margin=5*Plots.mm,
+    framestyle=:box
+)
+
+# Symbols encode sparsity levels; colors encode solvers.
+# Define marker for each sparsity and a color for each solver.
+## Sparsity shapes
+marker_map_sparsity = OrderedDict(0.1=>:circle, 0.5=>:square, 0.9=>:utriangle)
+## Solver color shades
+color_map_solver = OrderedDict("Default"=>:red, "gesv!"=>:blue, "Generated"=>:green)
+
+# Plot each point individually so marker shape shows sparsity and color shows solver.
+for solver_name in keys(solvers)
+    for sparsity in sparsity_levels
+        if sparsity in keys(results) && solver_name in keys(results[sparsity])
+            t = results[sparsity][solver_name].time
+            e = results[sparsity][solver_name].error
+            scatter!(p, [t], [e];
+                     label="",
+                     marker=marker_map_sparsity[sparsity],
+                     markersize=15,
+                     color=color_map_solver[solver_name],
+                     markerstrokecolor=:black,
+                     markerstrokewidth=0.0,#0.8,
+                     alpha=0.45)
+        end
+    end
+end
+
+# Create a combined legend
+for solver_name in keys(solvers)
+    for s in sparsity_levels
+        lbl = "$(solver_name), ρ:$(s)"
+        scatter!(p, [NaN], [NaN]; label=lbl,
+                 marker=marker_map_sparsity[s],
+                 markersize=15,
+                 color=color_map_solver[solver_name],
+                 markerstrokecolor=:black,
+                 markerstrokewidth=0.0, 
+                 alpha=0.45)
+    end
+end
+
+savefig(p, "error_vs_time.pdf")
+println("\n✓ Plot saved as error_vs_time.pdf")
+
+display(p)