Refactor Reproducer to Eliminate Code Duplication Using AST-Based Function Extraction

[FindHao]

fix #170

Summary

This PR eliminates duplicate functions between the reproducer template and utility module by implementing an automatic function extraction system using Python's AST (Abstract Syntax Tree) parser. The template file is reduced by 92% (398 → 31 lines) while maintaining full functionality.

Net Impact: 4 files changed, -40 lines of code (443 insertions, 483 deletions)

Problem

The reproducer system had significant code duplication:

• example.py (template) and utils.py contained 8 duplicate functions
• Maintaining consistency required updating code in two places
• utils.py was missing critical functionality:
  • ❌ No support for NoneType, str, float types
  • ❌ No stride/storage offset handling (critical for non-contiguous tensors)
  • ❌ No device normalization
  • ❌ No error handling for tensor capture failures

Solution

1. Single Source of Truth

All utility functions now live exclusively in utils.py. The template uses a placeholder that gets replaced with extracted code during reproducer generation.

2. AST-Based Function Extraction

Created function_extractor.py that uses Python's AST parser to:

• Extract functions and constants from source files without importing them
• Preserve original formatting, comments, and decorators
• Handle multi-line statements robustly

3. Enhanced Utility Functions

Fixed and expanded utils.py with:

• ✅ Support for all Python types (NoneType, str, float, etc.)
• ✅ Proper stride and storage offset handling for non-contiguous tensors
• ✅ Device normalization (cuda → cuda:0)
• ✅ Comprehensive error handling

Changes Overview

Files Modified

File	Lines Changed	Description
function_extractor.py	+122 (new)	AST-based extraction engine
placeholder_replacer.py	+11	Added handler for utility functions placeholder
example.py	-373	Removed all duplicate functions (92% reduction)
utils.py	+320/-	Enhanced with missing functionality

Detailed Changes

1. New Module: function_extractor.py (+122 lines)

Purpose: Extract utility functions from source files using AST parsing.

Key Functions:

def extract_utility_functions() -> str:
    """Main entry point - extracts all utility code"""
    
def _parse_source_file(file_path) -> tuple[ast.Module, list[str]]:
    """Parse Python file into AST and source lines"""
    
def _extract_assignment(tree, lines, var_name) -> str | None:
    """Extract module-level constants (e.g., TRITON_KERNELS_CUSTOM_TYPES)"""
    
def _extract_function(tree, lines, func_name) -> str | None:
    """Extract function definition including decorators"""
    
def _extract_functions(tree, lines, func_names) -> list[str]:
    """Batch extract multiple functions with error checking"""

Extracted Content:

• From utils.py: 8 functions + 1 constant
• From load_tensor.py: 1 function
• Total: ~14KB of code, 385 lines

Advantages over inspect-based approach:

• ✅ No module imports = no code execution = no side effects
• ✅ Unified extraction method for functions and constants
• ✅ Robust handling of decorators, multi-line statements, comments
• ✅ Uses official Python parser (not string manipulation)

2. Enhanced: utils.py (+320 lines, restructured)

Added Functions:

def create_args_from_json_file(json_path)
    """Load and parse JSON file"""

def _apply_stride_and_offset(tensor, shape, stride, storage_offset)
    """Apply custom stride and storage offset to tensors"""

def _create_base_tensor(arg_info)
    """Create base tensor without stride modifications"""

def _create_tensor(arg_info)
    """Create tensor with stride/offset applied"""

Enhanced Functions:

def create_args_from_json(data)
    # Refactored to accept parsed data instead of file path

def _create_arg_from_info(arg_info)
    # Added support for:
    # - NoneType, str, float types
    # - Stride and storage offset handling
    # - Device normalization (cuda → cuda:0)
    # - tensor_capture_error handling

Impact:

• Fixes critical missing functionality
• Enables proper handling of non-contiguous tensors
• Improves type coverage
• Better error handling

3. Simplified: example.py (-373 lines, -92%)

Before: 398 lines with duplicate function implementations

After: 31 lines with just structure and placeholders

"""
This file is automatically generated by TritonParse reproducer.
It contains a smallest testing example for a Triton kernel.
"""

import torch

# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}
# {{KERNEL_SYSPATH_PLACEHOLDER}}
# {{KERNEL_IMPORT_PLACEHOLDER}}
# {{UTILITY_FUNCTIONS_PLACEHOLDER}}  # <- NEW: Auto-injected utility code

if __name__ == "__main__":
    script_dir = Path(__file__).resolve().parent  # noqa: F821
    json_file = script_dir / "{{JSON_FILE_NAME_PLACEHOLDER}}"
    grid, args_dict = create_args_from_json_file(str(json_file))  # noqa: F821
    
    print("Generated kernel arguments dictionary:")
    for name, arg in args_dict.items():
        print(f"  {name}: {arg}")
    print(f"Grid: {grid}")
    
    # {{KERNEL_INVOCATION_PLACEHOLDER}}
    
    torch.cuda.synchronize()
    print("Kernel execution finished.")

Note: # noqa: F821 comments suppress linter warnings for identifiers that will be injected at generation time.

4. Updated: placeholder_replacer.py (+11 lines)

Changes:

# Added import
from tritonparse.reproducer.function_extractor import extract_utility_functions

# Added handler registration
self.register("# {{UTILITY_FUNCTIONS_PLACEHOLDER}}", self._replace_utility_functions)

# Added handler method
def _replace_utility_functions(self, code, context_bundle, **kwargs):
    """Replace the utility functions placeholder with extracted functions."""
    utility_code = extract_utility_functions()
    return code.replace("# {{UTILITY_FUNCTIONS_PLACEHOLDER}}", utility_code)

Technical Highlights

AST-Based Extraction

Why AST instead of inspect?

Aspect	inspect Module	AST Parser
Module Import	Required ⚠️	Not required ✅
Code Execution	Yes (side effects) ⚠️	No (static) ✅
Function Extraction	Simple ✅	Slightly more code ⚠️
Constant Extraction	String parsing hack ❌	Robust ✅
Consistency	Mixed approach ❌	Unified ✅
Extensibility	Limited ⚠️	Excellent ✅

Example: Robust Constant Extraction

Before (fragile string parsing):

# HACK: Loop through lines looking for pattern
constant_code = inspect.getsource(utils_module).split("\n")
for i, line in enumerate(constant_code):
    if line.startswith("TRITON_KERNELS_CUSTOM_TYPES"):
        j = i
        while not constant_code[j].endswith(")"):  # Assumes format!
            j += 1
        constant_lines = constant_code[i : j + 1]

After (robust AST-based):

# Use Python's official parser
for node in tree.body:
    if isinstance(node, ast.Assign):
        for target in node.targets:
            if isinstance(target, ast.Name) and target.id == var_name:
                # Parser provides exact line numbers
                return "\n".join(lines[node.lineno-1:node.end_lineno])

Workflow

How reproducer generation works now:

1. Load Template (example.py) - Only 31 lines with placeholders
2. Replace Placeholders:
   • {{JSON_FILE_NAME_PLACEHOLDER}} → JSON filename
   • {{KERNEL_IMPORT_PLACEHOLDER}} → Kernel import statement
   • {{UTILITY_FUNCTIONS_PLACEHOLDER}} → Extracted utility code (~14KB)
   • {{KERNEL_INVOCATION_PLACEHOLDER}} → Kernel call
3. Extract Functions - function_extractor.py parses source files via AST
4. Inject Code - Complete utility functions inserted into reproducer
5. Output - Standalone, executable Python script (no tritonparse dependency)

Testing & Validation

✅ All Tests Pass

1. Function Extraction Validation

✅ Extracts load_tensor function
✅ Extracts _get_triton_tensor_types function
✅ Extracts create_args_from_json_file function
✅ Preserves @lru_cache decorator
✅ Extracts TRITON_KERNELS_CUSTOM_TYPES constant
✅ Includes all necessary imports

2. Integration Test

$ python -m unittest tests.test_tritonparse.TestTritonparseCUDA.test_reproducer_end_to_end -v
test_reproducer_end_to_end ... ok
----------------------------------------------------------------------
Ran 1 test in 5.015s
OK

3. Code Quality

$ make format
✅ usort - Import sorting passed
✅ ruff - Linting passed
✅ black - Formatting passed

4. Output Consistency

• Before: 13,845 characters, 385 lines
• After: 13,835 characters, 376 lines
• ✅ Nearly identical output (minor whitespace differences expected)

Benefits

1. Code Maintainability 📈

Metric	Before	After	Change
Template file size	398 lines	31 lines	⬇️ 92%
Duplicate functions	8	0	✅ Eliminated
Single source of truth	❌	✅	✅ Achieved
Net lines of code	-	-	⬇️ -40 lines

2. Functional Completeness ✅

Feature	utils.py (before)	utils.py (after)
NoneType support	❌	✅
str/float support	❌	✅
Stride handling	❌	✅
Device normalization	❌	✅
Error handling	❌	✅

3. Advantages Over Previous Approach

• ✅ Single Source of Truth: All utility functions defined once in utils.py
• ✅ Easier Maintenance: Changes only needed in one place
• ✅ Automatic Sync: Generated reproducers always have latest implementations
• ✅ Simplified Template: example.py is cleaner and focused on structure
• ✅ No Code Execution: AST parsing doesn't execute code (safer)
• ✅ Enhanced Functionality: utils.py now feature-complete
• ✅ Future-Proof: Easy to add extraction for classes, type aliases, etc.

4. Developer Experience 🎯

• Before: Update function → Update in 2 places → Risk inconsistency
• After: Update function → Automatic propagation → Always consistent
• New feature: Add function to utils.py → Add name to extraction list → Done

Breaking Changes

None. This is a pure refactoring:

• ✅ Generated reproducers remain fully standalone (no tritonparse installation required)
• ✅ All existing functionality preserved
• ✅ All tests pass
• ✅ Public API unchanged
• ✅ Output format identical

Migration Notes

No action required from users. This change is completely transparent:

• Existing reproducers continue to work
• New reproducers generated with this change are functionally identical
• No changes to command-line interface or Python API

Future Possibilities

With AST-based extraction, we can now easily:

1. Extract Classes

   def _extract_class(tree, lines, class_name):
       # Easy to implement with same pattern

2. Extract Type Aliases

   def _extract_type_alias(tree, lines, alias_name):
       # Same AST-based approach

3. Selective Extraction

   • Extract only public functions (no _ prefix)
   • Filter by decorators
   • Analyze dependencies automatically

4. Cross-File Analysis

   • Detect unused functions
   • Find circular dependencies
   • Generate import graphs

Commits

1. Add function extractor module for reproducer utility functions

   • Initial implementation using inspect module
   • Mixed approach (inspect for functions, string parsing for constants)

2. Refactor function extractor to use AST parsing

   • Complete rewrite to pure AST-based extraction
   • Eliminates string parsing hacks
   • Unified extraction method for all code elements

Review Checklist

• All tests pass
• Code quality checks pass (usort, ruff, black)
• Generated reproducers tested and working
• No breaking changes
• Documentation updated
• Output validated for consistency

Conclusion

This PR successfully modernizes the reproducer system by:

• 🎯 Eliminating code duplication (8 functions)
• 🔧 Fixing critical missing functionality (stride handling, type support)
• 📉 Reducing template size by 92%
• 🛡️ Improving robustness (AST vs string parsing)
• 🔄 Enabling easier future maintenance
• ⚡ Maintaining full backward compatibility

The refactoring is complete, tested, and ready for review.

Test Plan

% python -m unittest tests.test_tritonparse -v -k test_reproducer_end_to_end
test_reproducer_end_to_end (tests.test_tritonparse.TestTritonparseCUDA.test_reproducer_end_to_end)
End-to-end test for reproducer: generate logs, build script, run it. ... Successfully converted to prettified JSON: /tmp/tmptgi9r4yq/repro_output/add_kernel/repro_context_20251020124238.json
INFO:tritonparse:REPRODUCER_OUTPUT
{'kernel_src_path': '', 'kernel': 'add_kernel', 'repro_script': '/tmp/tmptgi9r4yq/repro_output/add_kernel/repro_20251020124238.py', 'repro_context': '/tmp/tmptgi9r4yq/repro_output/add_kernel/repro_context_20251020124238.json'}
✓ Cleaned up temporary directory
ok

----------------------------------------------------------------------
Ran 1 test in 5.046s

OK

[meta-codesync]

@FindHao has imported this pull request. If you are a Meta employee, you can view this in D85093152.

[meta-codesync]

@FindHao merged this pull request in 89a72af.