optimize-code

name: optimize-code description: "Optimize a script for speed: parallelize loops, vectorize, reduce allocations, exploit GPU/SIMD, cache computations. Reads the input file, produces an optimized copy. Supports Julia, Python, R." argument-hint: " [output_file]" allowed-tools: ["Read", "Write", "Grep", "Glob", "Bash", "Agent"]

Optimize Code for Speed

Accept a source file, analyze it for performance bottlenecks, and produce an optimized version.

Arguments

• $ARGUMENTS format: <input_file> [output_file]
• If output_file is omitted, default to <input_file_stem>_optimized.<ext> in the same directory.

Workflow

Phase 1: Profile & Diagnose

Read the input file and identify the language (Julia, Python, R, etc.).

Build a bottleneck inventory by scanning for these patterns:

Julia-specific

Python-specific

R-specific

Phase 2: Optimization Strategy

For each bottleneck found, classify the fix:

1. Parallelize — @threads, Threads.@spawn, @distributed, pmap, joblib.Parallel
2. Vectorize — replace scalar loops with matrix/broadcast ops
3. Pre-allocate — replace push!/append with pre-sized arrays
4. Cache/Memoize — compute once, reuse (e.g., kernel matrices, sort indices)
5. GPU offload — move large matrix ops to CuArray/torch.cuda
6. Reduce allocation — @views, in-place operations (.=, mul!, ldiv!)
7. Type-stabilize — add type annotations, avoid Any containers
8. Algorithm change — better asymptotic complexity (e.g., avoid O(n²) when O(n log n) exists)

Rank fixes by expected speedup × implementation safety. Prioritize:

• Safe, high-impact changes first (pre-allocate, vectorize, cache)
• Parallel changes second (need to verify independence)
• Algorithm changes last (highest risk of introducing bugs)

Phase 3: Generate Optimized File

Write the optimized file to output_file. Follow these rules:

1. Preserve correctness — the optimized file MUST produce identical outputs to the original. When in doubt, keep the original.
2. Add comments — mark every optimization with # OPT: <description> so the user can review.
3. Keep structure — maintain the same function signatures, section organization, and variable names where possible. Don't refactor beyond what's needed for speed.
4. Benchmark hooks — add timing instrumentation around the optimized sections:
   • Julia: @time or @elapsed blocks
   • Python: time.perf_counter() pairs
   • R: system.time() blocks
5. Fallback safety — for GPU operations, keep CPU fallback. For threading, ensure thread-safe access.

Julia-specific optimizations to apply:

# Pre-allocate instead of push!
result = Vector{Float64}(undef, n)  # OPT: pre-allocate
for i in 1:n
    result[i] = compute(i)
end

# @threads for independent iterations
Threads.@threads for i in 1:n  # OPT: parallelize independent loop
    result[i] = expensive_fn(i)
end

# @views to avoid copies
@views X_sub = X[mask, :]  # OPT: view instead of copy

# @inbounds for safe indexed loops
@inbounds for i in 1:n  # OPT: skip bounds check (loop range verified)
    y[i] = X[i, :] .* w
end

# In-place operations
mul!(ZtZ, Z', Z)  # OPT: in-place matrix multiply
ldiv!(cholesky(A), b)  # OPT: in-place solve

# Cache repeated computations
# OPT: cache — computed once, used N times
cached_result = expensive_computation()
for i in 1:n
    use(cached_result)
end

# Type-stable containers
result = Dict{Int, Float64}()  # OPT: typed dict (was Dict())

# Broadcast fusion
@. y = a * x + b  # OPT: fused broadcast (single allocation)

Phase 4: Report

After writing the optimized file, print a summary:

Optimization Report: <input_file> → <output_file>
═══════════════════════════════════════════════════

Bottlenecks found: N
Optimizations applied: M

| # | Line(s) | Type | Description | Est. Speedup |
|---|---------|------|-------------|-------------|
| 1 | 150-180 | Parallelize | @threads on CD grid search | 8-16x |
| 2 | 320-340 | Pre-allocate | Replace push! with pre-sized Vector | 2-3x |
| ...

Unchanged sections: [list sections left as-is and why]

To verify correctness:
  1. Run original:   julia <input_file>
  2. Run optimized:  julia <output_file>
  3. Compare outputs: diff output/cv_results.csv output/cv_results_opt.csv

Phase 5: Correctness Guard

Before finishing, verify:

• All function signatures unchanged
• All output file paths unchanged (or parameterized)
• Random seed preserved in same position
• No data races in parallel sections (each thread writes to independent indices)
• GPU operations have CPU fallback
• No global mutable state accessed from @threads without locks

If any check fails, revert that specific optimization and note it in the report.