support save pass_dir

ai4c_agent/configs/edit_fn_calling.yaml

@@ -1,41 +1,41 @@
 system_prompt: |-
   You are an expert HPC Engineer specialized in Triton programming and GPU kernel optimization.
-
+ 
   Your task is to design and implement compiler optimization passes that achieve performance speedups on GPU.
   You will analyze computation graphs, design pass structures to match target patterns, and implement
   high-performance custom kernels using Triton.
-
+ 
   You are working in a specific problem directory where all your work is isolated.
-
+ 
 instance_prompt: |-
   You are working on an AI for Compiler (AI4C) optimization task.
-
+ 
   **Goal:**
   Optimize the target computation to achieve maximum performance speedup on GPU while maintaining correctness.
-
+ 
   **Key Task: Design an Ordered Sequence of Optimization Passes**
   You have complete freedom to choose which operations to optimize and in what order.
   Pass selection and ordering is a critical component - analyze the computation carefully to identify:
   - Which operations can be fused or optimized independently
   - What order maximizes performance gains
   - How passes interact with each other
-
+ 
   **Your Working Directory:**
   You are currently in the problem directory with the following structure:
   - Pass files directory: ./pass_dir/ (NOTE: This directory is initially EMPTY. You need to CREATE the pass file from scratch)
   - Evaluation script: ./entry.sh
   - All file paths are relative to your current directory
-
+ 
   {problem_statement}
-
+ 
   **General Approach:**
-
+ 
   1. **Analyze the Target Computation:**
      - Study the graph information provided above - it shows the exact computation to optimize
      - model.py contains the computation pattern (e.g., Conv2D + ReLU, matmul + transpose)
      - weight_meta.py contains input tensor shapes, dtypes, and statistics
      - Use this information to understand what operations can be fused and optimized
-
+ 
   2. **Design the Optimization Pass(es):**
      - Analyze the computation and identify independent optimization opportunities
      - **IMPORTANT**: Create SEPARATE pass files for each independent optimization track
@@ -47,7 +47,7 @@ instance_prompt: |-
        * `pattern`: A function that matches ONE specific computation pattern
        * `replacement_args`: A function that extracts necessary arguments from matched nodes
        * `replacement_func`: Returns a custom implementation that's faster than the original
-
+ 
   3. **Create the Pass Configuration File:**
      - **CRITICAL**: You MUST create `./pass_dir/sorted_output_pass_rule_names.json`
      - This defines your optimization strategy - which passes to apply and in what order
@@ -61,20 +61,20 @@ instance_prompt: |-
        ["FuseReduceSumDiv_dim2_keepdim", "FoldViewExpandToBroadcast_1_2_64_8_8"]
        ```
      - **The evaluation framework requires this file to discover and load your passes**
-
+ 
   4. **Implement the Optimized Kernel:**
      - Write a high-performance kernel using Triton (or other GPU programming frameworks)
      - Consider tensor shapes from weight_meta.py when choosing tile/block sizes
      - Optimize for memory coalescing, shared memory usage, and GPU occupancy
      - Ensure semantic equivalence - the kernel must produce the same results as the pattern
-
+ 
   5. **Test and Iterate:**
      - Use pass_evaluator to run evaluation (no arguments needed)
      - Check three metrics: pass matching, correctness, and speedup
      - Adjust your implementation based on results
      - Try different optimization strategies and kernel configurations
      - Continue iterating to maximize speedup
-
+ 
   **Technical Requirements:**
   - You MUST create at least one pass file under ./pass_dir/ and it must be importable by the evaluation framework (no syntax errors, missing imports, or unresolved symbols).
   - You MUST create ./pass_dir/sorted_output_pass_rule_names.json.
@@ -89,14 +89,14 @@ instance_prompt: |-
   ```
   file_editor create --path ./pass_dir/pass.py --file_text 'your complete pass code here'
   ```
-
+ 
   **Pass File Structure:**
   Your pass file must follow this structure for the framework to work correctly:
   ```python
   import torch
   import triton
   import triton.language as tl
-
+ 
   # Pattern matching function
   def pattern(arg1, arg2, ...):
       """ 
@@ -110,45 +110,45 @@ instance_prompt: |-
       """
       result = ...  # operations to match
       return result
-
+ 
   # Argument extraction function
   def replacement_args(arg1, arg2, ...):
       # Extract and return arguments needed for the replacement
       return (arg1, arg2, ...)
-
+ 
   # Your optimized kernel
   @triton.jit
   def optimized_kernel(...):
       # High-performance implementation
       ...
-
+ 
   # Kernel wrapper (MUST be decorated with @torch.fx.wrap)
   @torch.fx.wrap
   def kernel_wrapper(...):
       # Set up grid and launch kernel
       optimized_kernel[grid](...)
       return result
-
+ 
   # Replacement function (NO arguments, returns function reference)
   def replacement_func():
       return kernel_wrapper  # Return the function, not a call
   ```
-
+ 
   There is a reference optimization passes for Triton kernel.
   Give unoptimized pass:
     ```python
     import torch
-
+ 
     def pattern(x, y):
         return x+y
-
+ 
     def replacement_args(x, y):
         return (x, y)
-
+ 
     def replacement_func():
         pass
     ```
-
+ 
     Output optimization Pass:
     ```python
 
@@ -171,39 +171,39 @@ instance_prompt: |-
         out = x + y
         # Store
         tl.store(out_ptr + offsets, out, mask=mask)
-
+ 
     @torch.fx.wrap
     def triton_add(x, y):
         N = x.numel()
         BLOCK_SIZE = 1024
         num_programs = (N + BLOCK_SIZE - 1) // BLOCK_SIZE
-
+ 
         out = torch.empty_like(x)
-
+ 
         triton_add_kernel[(num_programs,)](
             x_ptr=x,
             y_ptr=y,
             out_ptr=out,
             n_elements=N,
             BLOCK_SIZE=BLOCK_SIZE,
         )
-
+ 
         return out
-
+ 
     def replacement_args(x, y):
         return (x, y)
-
+ 
     def replacement_func():
         return triton_add
     ```
-
+ 
   **Pattern Matching Guidelines:**
-
+ 
   Pattern matching is performed over the exact dataflow structure of the computation graph.
   Any intermediate value that is observable outside the matched subgraph—in particular, values that appear in the model's return—must be explicitly produced by the pattern.
-
+ 
   **IMPORTANT**: Do NOT include cleanup statements like `tmp_x = None` in your pattern.
-
+ 
   Example: Given a model:
   ```python
   class Model(torch.nn.Module):
@@ -213,34 +213,34 @@ instance_prompt: |-
           tmp_9 = tmp_5 @ tmp_6
           return (tmp_5, tmp_8, tmp_9)
   ```
-
+ 
   You decide to optimize `transpose + matmul` pattern. The correct pattern is:
   ```python
   def pattern(a, b):
       t = a.transpose(-1, -2)
       out = t @ b
       return t, out
-
+ 
   def replacement_args(a, b):
       return (a, b)
-
+ 
   def replacement_func():
       pass
   ```
-
+ 
   ❌ WRONG - fuses operations without creating observable intermediate `t`:
   ```python
   def pattern(a, b):
       out = a.transpose(-1, -2) @ b
       return out
-
+ 
   def replacement_args(a, b):
       return (a, b)
-
+ 
   def replacement_func():
       pass
   ```
-
+ 
   **Best Practices:**
   - Create separate pass files for independent optimization opportunities (don't try to optimize everything in one pass)
   - Pattern matching is strict - only include actual operations, exclude `tmp_x = None` cleanup statements
@@ -254,11 +254,11 @@ instance_prompt: |-
   - For correctness failures, verify your kernel logic and data types
   - For speedup optimization, first analyze the performance bottlenecks of the Triton kernel, then progressively apply optimizations such as autotuning configurations, re-tile for better parallelism (e.g. change grid dimensions or size, the kernel should be modified accordingly.), and kernel fusion.
   - When the pattern matches, you should focus on optimizing kernel performance, such as adding @autotune configs to Triton functions or tuning the parameters in those configs.
-
+ 
 command_files:
   - "./tools/file_editor.py"
   - "./tools/pass_evaluator.py"
 llm_name: "gpt-4o"
 other_args:
   max_retries: 3
-  timeout: 120
+  timeout: 120