⚡️ Speed up method InjectPerfOnly.collect_instance_variables by 769% in PR #1418 (fix/pytorch-forward-method-instrumentation)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1418

If you approve this dependent PR, these changes will be merged into the original PR branch fix/pytorch-forward-method-instrumentation.

This PR will be automatically closed if the original PR is merged.

---

📄 769% (7.69x) speedup for InjectPerfOnly.collect_instance_variables in codeflash/code_utils/instrument_existing_tests.py

⏱️ Runtime : 1.30 milliseconds → 150 microseconds (best of 15 runs)

📝 Explanation and details

The optimized code achieves a 768% speedup (from 1.30ms to 150μs) by replacing the expensive ast.walk() traversal with a targeted manual traversal strategy.

Key Optimization:

The original code uses ast.walk(func_node), which recursively visits every node in the entire AST tree - including all expression nodes, operators, literals, and other irrelevant node types. The line profiler shows this single loop consumed 87.3% of the execution time (9.2ms out of 10.5ms).

The optimized version implements a work-list algorithm that only traverses statement nodes (body, orelse, finalbody, handlers). This dramatically reduces the number of nodes examined:

• Original: 1,889 nodes visited per call
• Optimized: ~317 nodes visited per call (83% reduction)

Why This Works:

1. Targeted traversal: Assignment statements (ast.Assign) can only appear as statements, not as expressions buried deep in the tree. By only following statement-level structure (body, orelse, etc.), we skip visiting thousands of irrelevant expression nodes.

2. Cache-friendly: Local variables class_name and instance_vars eliminate repeated self. attribute lookups, reducing pointer indirection.

3. Early filtering: The manual stack-based approach allows us to skip entire branches of the AST that can't contain assignments.

Performance Impact by Test Case:

• Simple cases (single assignment): ~500-600% faster
• Complex nested cases: ~429% faster
• Large-scale scenario (300 assignments): 807% faster - showing the optimization scales particularly well with code complexity

The optimization preserves all functionality (same nodes discovered, same instance variables collected) while dramatically reducing the algorithmic complexity from O(all_nodes) to O(statement_nodes).

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 18 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

import ast  # use to parse source strings into AST nodes for testing
from types import \
    SimpleNamespace  # lightweight container for required attributes

import pytest  # used for our unit tests
from codeflash.code_utils.instrument_existing_tests import InjectPerfOnly

# NOTE: We avoid constructing a full FunctionToOptimize if its constructor is unknown.
# Instead we pass a SimpleNamespace with the attributes the InjectPerfOnly ctor reads:
# - function_name
# - parents (a sequence whose single element has attribute 'type' == "ClassDef")
# - top_level_parent_name
#
# This follows the principle of not redefining domain classes; SimpleNamespace is a
# standard library container used only to provide the minimal attribute interface needed.

def _get_function_node_from_source(source: str, func_name: str = "forward") -> ast.FunctionDef:
    """
    Helper: parse Python source and return the ast.FunctionDef for func_name.
    Asserts that the function exists and returns that node.
    """
    module = ast.parse(source)
    # find the function def with the requested name
    for node in module.body:
        if isinstance(node, ast.FunctionDef) and node.name == func_name:
            return node
    raise AssertionError(f"Function {func_name!r} not found in source")

def _make_injector(class_name: str | None, function_name: str = "forward") -> InjectPerfOnly:
    """
    Helper: create an InjectPerfOnly instance with the desired class_name and function_name.
    If class_name is None we create a Function-like object that results in injector.class_name
    remaining None. Otherwise create parents list of length 1 with type == "ClassDef" so that
    injector.class_name is set to top_level_parent_name.
    """
    if class_name is None:
        parents = []  # will prevent setting class_name inside InjectPerfOnly.__init__
        top_name = None
    else:
        # parent object must have a 'type' attribute equal to "ClassDef"
        parents = [SimpleNamespace(type="ClassDef")]
        top_name = class_name

    fake_function = SimpleNamespace(
        function_name=function_name,
        parents=parents,
        top_level_parent_name=top_name,
    )
    # call_positions and module_path are not used by collect_instance_variables; pass simple values
    return InjectPerfOnly(function=fake_function, module_path="mod.py", call_positions=[])

def test_no_collection_when_not_forward():
    # Scenario: Injector's only_function_name != "forward" -> no collection should occur.
    injector = _make_injector(class_name="AlexNet", function_name="not_forward")
    # prepare a function AST that would normally create an instance variable
    src = """
def not_forward(x):
    model = AlexNet(x)
"""
    func_node = _get_function_node_from_source(src, func_name="not_forward")
    # call the method under test
    injector.collect_instance_variables(func_node) # 679ns -> 740ns (8.24% slower)

def test_collect_single_assignment_basic():
    # Scenario: Basic positive case - simple assignment of the class -> variable should be tracked.
    injector = _make_injector(class_name="AlexNet", function_name="forward")
    src = """
def forward(x):
    model = AlexNet(x)
"""
    func_node = _get_function_node_from_source(src)
    injector.collect_instance_variables(func_node) # 17.5μs -> 2.96μs (491% faster)

def test_collect_multiple_targets_in_assignment():
    # Scenario: chained assignment "a = b = ClassName()" should collect all Name targets.
    injector = _make_injector(class_name="AlexNet", function_name="forward")
    src = """
def forward():
    a = b = AlexNet()
"""
    func_node = _get_function_node_from_source(src)
    injector.collect_instance_variables(func_node) # 15.4μs -> 2.56μs (502% faster)

def test_ignore_attribute_target_assignments():
    # Scenario: attribute assignment (e.g., self.model = ClassName()) should NOT be collected
    injector = _make_injector(class_name="AlexNet", function_name="forward")
    src = """
def forward():
    self.model = AlexNet()
"""
    func_node = _get_function_node_from_source(src)
    injector.collect_instance_variables(func_node) # 15.0μs -> 2.23μs (573% faster)

def test_ignore_non_matching_class_name():
    # Scenario: call to a different class with same pattern should not be collected
    injector = _make_injector(class_name="AlexNet", function_name="forward")
    src = """
def forward():
    candidate = OtherNet()
"""
    func_node = _get_function_node_from_source(src)
    injector.collect_instance_variables(func_node) # 12.8μs -> 1.96μs (554% faster)

def test_ignore_calls_where_func_is_attribute():
    # Scenario: calls like models.AlexNet() have node.value.func as ast.Attribute, not ast.Name,
    # and therefore should not be collected by the current implementation.
    injector = _make_injector(class_name="AlexNet", function_name="forward")
    src = """
def forward():
    m = models.AlexNet()
"""
    func_node = _get_function_node_from_source(src)
    injector.collect_instance_variables(func_node) # 14.3μs -> 2.02μs (607% faster)

def test_tuple_assignment_not_collected():
    # Scenario: tuple targets (a, b) = ClassName() -> top-level target is ast.Tuple, implementation
    # only accepts ast.Name targets so no names inside tuples are collected.
    injector = _make_injector(class_name="AlexNet", function_name="forward")
    src = """
def forward():
    (a, b) = AlexNet()
"""
    func_node = _get_function_node_from_source(src)
    injector.collect_instance_variables(func_node) # 15.9μs -> 2.40μs (560% faster)

def test_nested_assignment_inside_blocks_is_collected():
    # Scenario: an assignment nested inside an if block should still be discovered via ast.walk.
    injector = _make_injector(class_name="AlexNet", function_name="forward")
    src = """
def forward(cond):
    if cond:
        inside = AlexNet()
    else:
        other = AlexNet()
"""
    func_node = _get_function_node_from_source(src)
    injector.collect_instance_variables(func_node) # 20.3μs -> 3.84μs (429% faster)

def test_large_scale_many_assignments():
    # Large scale scenario: create many assignments but keep well under the 1000-element guideline.
    # This verifies the collector scales to many targets and remains correct.
    count = 300  # well under 1000 to respect test constraints
    class_name = "C"
    injector = _make_injector(class_name=class_name, function_name="forward")
    # Build a function source with many assignments: v0 = C(); v1 = C(); ...
    lines = ["def forward():"] + [f"    v{i} = {class_name}()" for i in range(count)]
    src = "\n".join(lines)
    func_node = _get_function_node_from_source(src)
    injector.collect_instance_variables(func_node) # 1.19ms -> 131μs (807% faster)
    # Expect all variable names to be collected
    expected = {f"v{i}" for i in range(count)}
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-pr1418-2026-02-06T22.39.42 and push.