Starting new dagger example

emmanuellujan · emmanuellujan · commit 0f1b4ed1571d · 2025-12-05T17:20:23.000-05:00
diff --git a/Project.toml b/Project.toml
@@ -9,9 +9,11 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DecisionTree = "7806a523-6efd-50cb-b5f6-3fa6f1930dbb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 MKL = "33e6dc65-8f57-5167-99aa-e5a354878fb2"
 MatrixDepot = "b51810bb-c9f3-55da-ae3c-350fc1fbce05"
 OpenAI = "e9f21f70-7185-4079-aca2-91159181367c"
@@ -21,15 +23,20 @@ ScikitLearn = "3646fa90-6ef7-5e7e-9f22-8aca16db6324"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
+[sources]
+Dagger = {rev = "master", url = "https://github.com/JuliaParallel/Dagger.jl"}
+
 [compat]
+BSON = "0.3"
 BenchmarkTools = "1"
 CSV = "0.10"
-BSON = "0.3"
 CUDA = "5.9.2"
 CairoMakie = "0.15"
-DecisionTree = "0.12"
+Dagger = "0.19.2"
 DataFrames = "1"
+DecisionTree = "0.12"
 LinearAlgebra = "1.12.0"
+Logging = "1.11.0"
 MKL = "0.9"
 MatrixDepot = "1.0.13"
 OpenAI = "0.12.0"
diff --git a/examples/agentic/generate-dagger-linear-solver/Project.toml b/examples/agentic/generate-dagger-linear-solver/Project.toml
@@ -0,0 +1,12 @@
+[deps]
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+SmartSolve = "4fbb3a3c-2fa1-4c19-8d57-bae8bc1e16ac"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+
+[sources]
+Dagger = {rev = "master", url = "https://github.com/JuliaParallel/Dagger.jl"}
+SmartSolve = {path = "../../.."}
diff --git a/examples/agentic/generate-dagger-linear-solver/benchmark.jl b/examples/agentic/generate-dagger-linear-solver/benchmark.jl
@@ -0,0 +1,136 @@
+using LinearAlgebra
+using SparseArrays
+using CUDA
+using BenchmarkTools
+using OrderedCollections
+using Plots
+
+println("GPU benchmark with error-vs-time plot:\n")
+
+include("solver.jl")
+
+# Configuration
+N = 15_000
+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) / 1e9  # Convert to s
+            
+            # 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)) s, 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 (s)",
+    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,
+    title="Random Matrices of Size $(N)x$(N),\n Varying Sparsity Levels (ρ) and\n GPU Solvers",
+    titlefontsize=22,
+)
+
+# 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)
diff --git a/examples/agentic/generate-dagger-linear-solver/generate.jl b/examples/agentic/generate-dagger-linear-solver/generate.jl
@@ -0,0 +1,19 @@
+using SmartSolve
+using LinearAlgebra
+using SparseArrays
+using CUDA
+using BenchmarkTools
+using Dagger
+
+prompt = """
+Generate a high-performance Dagger.jl (https://juliaparallel.org/Dagger.jl/dev/) implementation in Julia of a linear solver for sparse matrices
+based on LU with iterative refinement (at least 5 refinement iterations), using the following
+reference: https://nhigham.com/2023/03/13/what-is-iterative-refinement
+"""
+
+secret_key = ENV["OPENAI_API_KEY"]
+solver, hist, conv = gen_linear_solver_dagger(prompt, secret_key; max_iters = 5)
+
+println("Generated Code:\n")
+println(solver)
+write("solver.jl", solver)
diff --git a/examples/agentic/generate-dagger-linear-solver/readme b/examples/agentic/generate-dagger-linear-solver/readme
@@ -0,0 +1 @@
+This example generates a high-performance Dagger.jl implementation in Julia for solving sparse linear systems using an LU-based method with iterative refinement.
diff --git a/src/Agentic.jl b/src/Agentic.jl
@@ -20,7 +20,9 @@ function error_prompt_maker(err_message)
 end
 
 proposed_fn(x) = x
-function generate_default_code(prompt, secret_key, checker_filename, model = "gpt-5-mini", dev_prompt_fn = dev_prompt_maker; max_iters = 3)
+evaluator(x) = (true, "")
+function generate_default_code(prompt, secret_key, checker_filename;
+                              model = "gpt-5-mini", dev_prompt_fn = dev_prompt_maker, max_iters = 3)
     """
         - checker_fn: proposed_fn -> check : Bool, performance_description : String
     """
@@ -81,18 +83,64 @@ function ls_cuda_dev_prompt_maker(fn_str)
             " Assume that LinearAlgebra and SparseArrays is already imported."
 end
 
+function ls_dagger_dev_prompt_maker(fn_str)
+    return "You are a numerical linear algebra expert, and an expert Julia programmer. You are very experienced in GPU programming using CUDA." * 
+            " The user will ask you to generate a function and use the following code the check if your solution is accurate and fast." * 
+            " Make sure the code you produce uses Dagger." *
+            " Here is the code: \n" * fn_str * "\nOnly return the function. Make sure the function name is proposed_fn. Do not return extra text." *
+            " Assume that LinearAlgebra and SparseArrays is already imported." *
+            " Assume that Dagger is already imported." *
+            " Use the following Dagger.jl documentation: https://juliaparallel.org/Dagger.jl/dev/" *
+            " Use the following Dagger.jl implementation of Cholesky as an example: https://github.com/JuliaParallel/Dagger.jl/blob/67211816781d59109d74940550ca2d80af96b13d/src/array/cholesky.jl"
+end
+
 src_dir = @__DIR__
 
-function gen_linear_solver(prompt, secret_key, checker_filename = src_dir * "/test_performance.jl", model = "gpt-5-mini"; max_iters = 10)
-    return generate_default_code(prompt, secret_key, checker_filename, model, ls_dev_prompt_maker; max_iters = max_iters)
+function gen_linear_solver(prompt, secret_key;
+                          checker_filename = src_dir * "/test_performance.jl",
+                          model = "gpt-5-mini",
+                          max_iters = 10)
+    return generate_default_code(prompt, secret_key, checker_filename;
+                                model = model,
+                                dev_prompt_fn=ls_dev_prompt_maker,
+                                max_iters = max_iters)
 end
 
-function gen_linear_solver_cuda(prompt, secret_key,checker_filename = src_dir *"/test_performance_cuda.jl", model = "gpt-5-mini"; max_iters = 10)
-    return generate_default_code(prompt, secret_key, checker_filename, model, ls_cuda_dev_prompt_maker; max_iters = max_iters)
+function gen_linear_solver_cuda(prompt, secret_key;
+                                checker_filename = src_dir *"/test_performance_cuda.jl",
+                                model = "gpt-5-mini",
+                                max_iters = 10)
+    return generate_default_code(prompt, secret_key, checker_filename;
+                                 model=model,
+                                 dev_prompt_fn=ls_cuda_dev_prompt_maker,
+                                 max_iters = max_iters)
+end
+
+function gen_linear_solver_dagger(prompt, secret_key;
+                                  checker_filename = src_dir *"/test_performance_dagger.jl",
+                                  model = "gpt-5-mini",
+                                  max_iters = 10)
+    return generate_default_code(prompt, secret_key, checker_filename;
+                                 model=model,
+                                 dev_prompt_fn=ls_dagger_dev_prompt_maker,
+                                 max_iters = max_iters)
 end
 
 function printhist(hist)
     for (i, (role, message)) in enumerate(hist)
             println("Message $i $(role[2]):\n$(message[2])\n")
     end
+end
+
+function get_report(m_err, m_runtime, m_alloc,
+                    err_threshold, runtime_threshold, alloc_threshold)
+    report = """
+    Median error ratio (error_default / error_gen): $(m_err)
+    Desired median error ratio: >= $err_threshold
+    Median runtime ratio or speedup (runtime_default / runtime_gen): $(m_runtime)
+    Desired median runtime ratio: >= $runtime_threshold
+    Allocation median ratio (alloc_default / alloc_gen): $(m_alloc)
+    Desired median allocation ratio: >= $alloc_threshold
+    """
+    return report
 end
diff --git a/src/SmartSolve.jl b/src/SmartSolve.jl
@@ -1,5 +1,6 @@
 module SmartSolve
 
+
 using MatrixDepot
 using LinearAlgebra
 using DataFrames
@@ -13,15 +14,14 @@ using BSON
 using SparseArrays
 using OpenAI
 using CUDA
+using Dagger
 
 include("SmartDiscovery.jl")
 include("SmartDB.jl")
 include("SmartModel.jl")
 include("Utils.jl")
 include("Agentic.jl")
-include("test_performance.jl")
-# include("test_performance_cuda.jl")
 
-export generate_default_code, gen_linear_solver, gen_linear_solver_cuda, printhist
+export generate_default_code, gen_linear_solver, gen_linear_solver_cuda, gen_linear_solver_dagger, printhist
 
 end # module SmartSolve
diff --git a/src/test_performance.jl b/src/test_performance.jl
@@ -4,19 +4,6 @@ push!(test_matrices, sprand(N, N, 0.1))
 push!(test_matrices, sprand(N, N, 0.2))
 push!(test_matrices, sprand(N, N, 0.3))
 
-function get_report(m_err, m_runtime, m_alloc,
-                    err_threshold, runtime_threshold, alloc_threshold)
-    report = """
-    Median error ratio (error_default / error_gen): $(m_err)
-    Desired median error ratio: >= $err_threshold
-    Median Runtime ratio or speedup (runtime_default / runtime_gen): $(m_runtime)
-    Desired median runtime ratio: >= $runtime_threshold
-    Allocation median ratio (alloc_default / alloc_gen): $(m_alloc)
-    Desired median allocation ratio: >= $alloc_threshold
-    """
-    return report
-end
-
 function evaluator(proposed_fn, err_threshold=1.0,
                                 runtime_threshold=1.1,
                                 alloc_threshold=0.0)
diff --git a/src/test_performance_cuda.jl b/src/test_performance_cuda.jl
@@ -4,23 +4,10 @@ push!(test_matrices, sprand(N, N, 0.1))
 push!(test_matrices, sprand(N, N, 0.2))
 push!(test_matrices, sprand(N, N, 0.3))
 
-function get_report(m_err, m_runtime, m_alloc,
-                    err_threshold, runtime_threshold, alloc_threshold)
-    report = """
-    Median error ratio (error_default / error_gen): $(m_err)
-    Desired median error ratio: >= $err_threshold
-    Median runtime ratio or speedup (runtime_default / runtime_gen): $(m_runtime)
-    Desired median runtime ratio: >= $runtime_threshold
-    Allocation median ratio (alloc_default / alloc_gen): $(m_alloc)
-    Desired median allocation ratio: >= $alloc_threshold
-    """
-    return report
-end
-
-function evaluator_cuda(proposed_fn;
-                        err_threshold::Float64 = 1.0,
-                        runtime_threshold::Float64 = 1.1,
-                        alloc_threshold::Float64 = 0.0)
+function evaluator( proposed_fn;
+                    err_threshold::Float64 = 1.0,
+                    runtime_threshold::Float64 = 1.1,
+                    alloc_threshold::Float64 = 0.0)
 
     error_ratios  = Float64[]
     runtime_ratios = Float64[]
diff --git a/src/test_performance_dagger.jl b/src/test_performance_dagger.jl

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+This example generates a high-performance Dagger.jl implementation in Julia for solving sparse linear systems using an LU-based method with iterative refinement.`