⚡️ Speed up method JitDecoratorDetector._is_jit_decorator by 16% in PR #1048 (better-jit-detection)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1048

If you approve this dependent PR, these changes will be merged into the original PR branch better-jit-detection.

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

---

📄 16% (0.16x) speedup for JitDecoratorDetector._is_jit_decorator in codeflash/code_utils/line_profile_utils.py

⏱️ Runtime : 411 microseconds → 355 microseconds (best of 178 runs)

📝 Explanation and details

The optimized code achieves a 15% speedup by eliminating function call overhead and reducing redundant list operations in the hot path (_check_attribute_decorator).

Key Optimizations

1. Inlined Attribute Chain Extraction (Biggest Impact)

The original code called _get_attribute_parts(node) which incurred function call overhead and created a full list with reversal. The optimized version inlines this logic directly in _check_attribute_decorator, building the parts list in reverse order from the start. This eliminates:

• Function call overhead (~44% of original method time)
• One list reversal operation
• Early exit opportunities when current is not a Name

Performance impact: The line profiler shows _check_attribute_decorator improved from 2.23ms to 2.05ms (8% faster), with the inlined attribute extraction taking only 3.8% of method time vs the original 44.2% spent in the function call.

2. Dictionary Lookup Caching

Both _check_name_decorator and _check_attribute_decorator now use .get() to cache the alias lookup result instead of doing in check followed by dictionary access. This reduces two dictionary operations to one:

# Original: 2 dict operations
if name not in self.import_aliases:  # lookup 1
    return False
module, imported_name = self.import_aliases[name]  # lookup 2

# Optimized: 1 dict operation
alias_info = self.import_aliases.get(name)  # single lookup
if alias_info is None:
    return False
module, imported_name = alias_info

3. Streamlined List Operations

The original code created rest_parts = parts[1:] (a list slice) and then accessed rest_parts[-1] and rest_parts[:-1] multiple times. The optimized version:

• Keeps parts in reverse order (decorator name at index 0)
• Accesses decorator_name = parts[0] directly
• Only reverses sub-slices when building module paths (reversed(parts[1:]))

This reduces memory allocations and improves cache locality.

Test Case Performance

The optimization performs best on:

• Attribute decorator tests with module aliases (e.g., @nb.jit): 31.6% faster
• Unknown decorator checks: 20-32% faster due to early exit optimization
• Large-scale tests with many aliases: 19.7% faster showing the optimization scales well

Some simple name decorator cases are slightly slower (2-6%) due to the additional .get() overhead, but these are not the dominant workload based on the profiler showing attribute decorators consume 79.4% of total time.

Impact on Workloads

This optimization is particularly valuable for codebases that:

• Use attribute-based JIT decorators extensively (e.g., @numba.jit, @torch.jit.script)
• Have many import aliases that need to be resolved
• Perform static analysis on large codebases where this detector runs frequently

The 15% overall speedup comes primarily from optimizing the hot path in _check_attribute_decorator, which handles the most common and expensive decorator patterns.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import ast

# Import the real class under test from the provided module path.
from codeflash.code_utils.line_profile_utils import JitDecoratorDetector


# Basic tests
def test_simple_name_and_call_decorator_detection():
    """Basic functionality:
    - A simple name decorator that was imported via `from module import name` should be
      detected when the imported name maps directly to a known JIT decorator.
    - A Call node wrapping that Name (e.g., @jit()) should be detected via recursion.
    """
    det = JitDecoratorDetector()
    # Simulate: from tensorflow import function as tf_function_alias
    det.import_aliases = {"tf_function_alias": ("tensorflow", "function")}

    # Name node: @tf_function_alias -> should be True because tensorflow.function is a JIT decorator
    name_node = ast.Name(id="tf_function_alias")
    codeflash_output = det._is_jit_decorator(name_node)  # 1.32μs -> 1.35μs (2.22% slower)

    # Call node: @tf_function_alias() -> should be True because we unwrap Call and check its func
    call_node = ast.Call(func=ast.Name(id="tf_function_alias"), args=[], keywords=[])
    codeflash_output = det._is_jit_decorator(call_node)  # 751ns -> 801ns (6.24% slower)

    # A completely unknown name should be False
    unknown_node = ast.Name(id="not_in_aliases")
    codeflash_output = det._is_jit_decorator(unknown_node)  # 451ns -> 481ns (6.24% slower)


def test_attribute_detection_with_module_alias_and_direct_module_paths():
    """Attribute and module path cases:
    - import numba as nb; @nb.jit should be detected (module alias + attribute).
    - Direct module chain like @numba.cuda.jit should be detected even if no import alias is present.
    - Attribute with a non-name base (e.g., a Call as base) should safely return False.
    """
    det = JitDecoratorDetector()

    # Case 1: module alias that was imported (import numba as nb)
    det.import_aliases = {"nb": ("numba", None)}
    nb_jit_node = ast.Attribute(value=ast.Name(id="nb"), attr="jit")
    codeflash_output = det._is_jit_decorator(nb_jit_node)  # 2.79μs -> 2.12μs (31.6% faster)

    # Case 2: direct attribute chain without any alias: numba.cuda.jit
    # Build nested attribute AST: (numba).cuda.jit
    numba_name = ast.Name(id="numba")
    numba_cuda = ast.Attribute(value=numba_name, attr="cuda")
    numba_cuda_jit = ast.Attribute(value=numba_cuda, attr="jit")
    codeflash_output = det._is_jit_decorator(numba_cuda_jit)  # 2.48μs -> 2.83μs (12.1% slower)

    # Case 3: attribute whose value is a Call (not a Name). _get_attribute_parts should return []
    # e.g., @somefunc().jit  -> base isn't a Name; should return False
    call_base = ast.Call(func=ast.Name(id="somefunc"), args=[], keywords=[])
    attr_on_call = ast.Attribute(value=call_base, attr="jit")
    codeflash_output = det._is_jit_decorator(attr_on_call)  # 862ns -> 761ns (13.3% faster)


def test_from_import_attribute_access_and_name_behaviour():
    """Tests highlighting the nuance in the implementation:
    - from torch import jit; using @jit alone should NOT be detected because the detector expects
      attribute access to build 'torch.jit', but @jit (simple name) resolves to checking
      'torch' module for 'jit' which is not how JIT_DECORATORS is organized for torch.jit.
    - However, using @jit.script (i.e., attribute on the imported object) should be detected.
    """
    det = JitDecoratorDetector()

    # Simulate: from torch import jit
    det.import_aliases = {"jit": ("torch", "jit")}

    # Simple name @jit -> should be False based on implementation nuances (needs attribute to form torch.jit)
    simple_name = ast.Name(id="jit")
    codeflash_output = det._is_jit_decorator(simple_name)  # 1.34μs -> 1.36μs (1.40% slower)

    # Attribute @jit.script -> should be True (forms torch.jit.script)
    attr_node = ast.Attribute(value=ast.Name(id="jit"), attr="script")
    codeflash_output = det._is_jit_decorator(attr_node)  # 2.56μs -> 2.03μs (26.1% faster)


# Edge cases
def test_name_imported_module_used_directly_and_unknown_decorators():
    """Edge cases:
    - When a name maps to (module, None) (i.e., a module import alias), using the name directly
      as a decorator (e.g., @nb) should return False.
    - When the module is known but the decorator name is not in the module's JIT_DECORATORS,
      detection should return False.
    """
    det = JitDecoratorDetector()

    # Simulate: import numba as nb  (module import)
    det.import_aliases = {"nb": ("numba", None)}
    name_node = ast.Name(id="nb")
    codeflash_output = det._is_jit_decorator(name_node)  # 962ns -> 901ns (6.77% faster)

    # Simulate attribute numba.not_a_jit -> should be False (decorator name not in JIT_DECORATORS['numba'])
    not_jit_attr = ast.Attribute(value=ast.Name(id="numba"), attr="not_a_jit")
    codeflash_output = det._is_jit_decorator(not_jit_attr)  # 2.27μs -> 1.72μs (32.0% faster)


def test_attribute_chain_with_from_import_and_long_rest_parts():
    """More nuanced attribute reconstruction:
    - Simulate 'from somepkg import foo' and then use @foo.bar.baz
      The code should reconstruct a module path 'somepkg.foo.bar' and check the final attribute 'baz'.
    - We'll set up an alias mapping that matches such a reconstruction to a known JIT decorator.
    """
    det = JitDecoratorDetector()

    # Build a synthetic mapping such that: first_part -> ('pkg', 'subpkg')
    # We'll create a mapping and an attribute chain foo.bar.script -> should map to 'pkg.subpkg.bar', decorator 'script'
    # To make this resolve to a known decorator, register in JIT_DECORATORS a module 'pkg.subpkg.bar' with 'script'
    # Since tests must not modify JIT_DECORATORS from the source code (to respect "import real classes"),
    # we will instead choose a realistic mapping using existing JIT_DECORATORS entries:
    # For example, from torch import jit as tjit -> @tjit.trace should map to torch.jit.trace
    det.import_aliases = {"tjit": ("torch", "jit")}
    attr_node = ast.Attribute(value=ast.Attribute(value=ast.Name(id="tjit"), attr="trace"), attr="unused")
    # Here rest_parts will be ['trace', 'unused'] -> detector will form full_module = 'torch.jit.trace' and check decorator 'unused'
    # This decorator won't be known -> should be False
    codeflash_output = det._is_jit_decorator(attr_node)  # 3.49μs -> 3.43μs (1.78% faster)

    # Now construct @tjit.trace (i.e., only one rest part) which should map to torch.jit.trace (known)
    simple_attr = ast.Attribute(value=ast.Name(id="tjit"), attr="trace")
    codeflash_output = det._is_jit_decorator(simple_attr)  # 1.69μs -> 1.56μs (8.45% faster)


def test_attribute_base_not_name_returns_false():
    """Ensure that attribute chains whose base is not a Name node are rejected with False.
    This covers cases where AST attribute chains are malformed for our use (e.g., base is a literal).
    """
    det = JitDecoratorDetector()

    # base is a constant; @'numba'.jit doesn't make sense but test robustness
    const_base = ast.Constant(value="numba")
    attr_node = ast.Attribute(value=const_base, attr="jit")
    codeflash_output = det._is_jit_decorator(attr_node)  # 1.36μs -> 1.21μs (12.5% faster)


# Large scale tests (stress but bounded)
def test_large_scale_many_aliases_attribute_checks():
    """Large scale scenario:
    - Create a sizeable number of module alias entries (well under 1000 per instructions).
    - Verify that attribute lookups using each alias are resolved correctly and efficiently.
    - We keep the number modest (e.g., 300) to evaluate scalability without exhausting resources.
    """
    det = JitDecoratorDetector()

    # Create many aliases like nb0, nb1, ... mapping to the same module 'numba' (module import alias)
    num_aliases = 300  # well under the 1000 limit
    det.import_aliases = {}
    nodes = []
    for i in range(num_aliases):
        alias = f"nb{i}"
        det.import_aliases[alias] = ("numba", None)  # import numba as nb{i}
        # Build AST node nb{i}.jit
        nodes.append(ast.Attribute(value=ast.Name(id=alias), attr="jit"))

    # All should resolve to True because numba.jit is a known JIT decorator
    # Also assert that the detector doesn't accidentally mutate state across calls
    for node in nodes:
        codeflash_output = det._is_jit_decorator(node)  # 280μs -> 234μs (19.7% faster)


# Additional coverage: direct module chains and negative cases
def test_direct_module_chains_and_negative_examples():
    """Additional tests to cover:
    - Direct module chains like torch.jit.script and torch.compile
    - Ensure negative example torch.something_unknown is False
    """
    det = JitDecoratorDetector()
    det.import_aliases = {}  # no aliases - use direct module chains

    # torch.jit.script -> should be True (torch.jit.module 'script' exists)
    torch_name = ast.Name(id="torch")
    torch_jit = ast.Attribute(value=torch_name, attr="jit")
    torch_jit_script = ast.Attribute(value=torch_jit, attr="script")
    codeflash_output = det._is_jit_decorator(torch_jit_script)  # 3.58μs -> 3.62μs (1.08% slower)

    # torch.compile is a direct decorator under 'torch' module
    torch_compile = ast.Attribute(value=ast.Name(id="torch"), attr="compile")
    codeflash_output = det._is_jit_decorator(torch_compile)  # 1.51μs -> 1.16μs (30.2% faster)

    # torch.unknown should be False
    torch_unknown = ast.Attribute(value=ast.Name(id="torch"), attr="unknown_decorator")
    codeflash_output = det._is_jit_decorator(torch_unknown)  # 1.06μs -> 882ns (20.4% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import ast

import pytest
from codeflash.code_utils.line_profile_utils import JitDecoratorDetector


class TestJitDecoratorDetectorBasic:
    """Basic test cases for JitDecoratorDetector._is_jit_decorator"""

    def test_simple_name_jit_decorator(self):
        """Test detection of simple @jit decorator when jit is imported directly"""
        detector = JitDecoratorDetector()
        detector.import_aliases["jit"] = ("numba", "jit")
        
        # Parse @jit as a decorator
        code = "@jit\ndef func(): pass"
        tree = ast.parse(code)
        decorator_node = tree.body[0].decorator_list[0]
        
        # Should detect jit as a JIT decorator
        codeflash_output = detector._is_jit_decorator(decorator_node)

    def test_simple_name_non_jit_decorator(self):
        """Test that non-JIT decorators are not detected"""
        detector = JitDecoratorDetector()
        detector.import_aliases["cached"] = ("functools", "lru_cache")
        
        # Parse @cached as a decorator
        code = "@cached\ndef func(): pass"
        tree = ast.parse(code)
        decorator_node = tree.body[0].decorator_list[0]
        
        # Should not detect as JIT decorator
        codeflash_output = detector._is_jit_decorator(decorator_node)

    def test_attribute_decorator_numba_jit(self):
        """Test detection of @numba.jit style decorator"""
        detector = JitDecoratorDetector()
        detector.import_aliases["numba"] = ("numba", None)
        
        # Parse @numba.jit as a decorator
        code = "@numba.jit\ndef func(): pass"
        tree = ast.parse(code)
        decorator_node = tree.body[0].decorator_list[0]
        
        # Should detect numba.jit as a JIT decorator
        codeflash_output = detector._is_jit_decorator(decorator_node)

    def test_attribute_decorator_torch_jit_script(self):
        """Test detection of @torch.jit.script style decorator"""
        detector = JitDecoratorDetector()
        detector.import_aliases["torch"] = ("torch", None)
        
        # Parse @torch.jit.script as a decorator
        code = "@torch.jit.script\ndef func(): pass"
        tree = ast.parse(code)
        decorator_node = tree.body[0].decorator_list[0]
        
        # Should detect torch.jit.script as a JIT decorator
        codeflash_output = detector._is_jit_decorator(decorator_node)

    def test_call_decorator_with_args(self):
        """Test detection of decorator with call arguments @jit()"""
        detector = JitDecoratorDetector()
        detector.import_aliases["jit"] = ("numba", "jit")
        
        # Parse @jit() as a decorator with call
        code = "@jit()\ndef func(): pass"
        tree = ast.parse(code)
        decorator_node = tree.body[0].decorator_list[0]
        
        # Should detect jit() as a JIT decorator
        codeflash_output = detector._is_jit_decorator(decorator_node)

    def test_call_decorator_with_parameters(self):
        """Test detection of decorator with multiple parameters @jit(nopython=True)"""
        detector = JitDecoratorDetector()
        detector.import_aliases["jit"] = ("numba", "jit")
        
        # Parse @jit(nopython=True) as a decorator
        code = "@jit(nopython=True)\ndef func(): pass"
        tree = ast.parse(code)
        decorator_node = tree.body[0].decorator_list[0]
        
        # Should still detect jit() as a JIT decorator
        codeflash_output = detector._is_jit_decorator(decorator_node)

    def test_alias_numba_as_nb(self):
        """Test detection with import alias: import numba as nb"""
        detector = JitDecoratorDetector()
        detector.import_aliases["nb"] = ("numba", None)
        
        # Parse @nb.jit as a decorator
        code = "@nb.jit\ndef func(): pass"
        tree = ast.parse(code)
        decorator_node = tree.body[0].decorator_list[0]
        
        # Should detect nb.jit as a JIT decorator
        codeflash_output = detector._is_jit_decorator(decorator_node)

    def test_alias_njit_direct_import(self):
        """Test detection of njit when imported directly"""
        detector = JitDecoratorDetector()
        detector.import_aliases["njit"] = ("numba", "njit")
        
        # Parse @njit as a decorator
        code = "@njit\ndef func(): pass"
        tree = ast.parse(code)
        decorator_node = tree.body[0].decorator_list[0]
        
        # Should detect njit as a JIT decorator
        codeflash_output = detector._is_jit_decorator(decorator_node)

    def test_tensorflow_function_decorator(self):
        """Test detection of @tf.function decorator"""
        detector = JitDecoratorDetector()
        detector.import_aliases["tf"] = ("tf", None)
        
        # Parse @tf.function as a decorator
        code = "@tf.function\ndef func(): pass"
        tree = ast.parse(code)
        decorator_node = tree.body[0].decorator_list[0]
        
        # Should detect tf.function as a JIT decorator
        codeflash_output = detector._is_jit_decorator(decorator_node)

    def test_jax_jit_decorator(self):
        """Test detection of @jax.jit decorator"""
        detector = JitDecoratorDetector()
        detector.import_aliases["jax"] = ("jax", None)
        
        # Parse @jax.jit as a decorator
        code = "@jax.jit\ndef func(): pass"
        tree = ast.parse(code)
        decorator_node = tree.body[0].decorator_list[0]
        
        # Should detect jax.jit as a JIT decorator
        codeflash_output = detector._is_jit_decorator(decorator_node)


class TestJitDecoratorDetectorEdgeCases:
    """Edge case tests for JitDecoratorDetector._is_jit_decorator"""

    def test_undefined_name_decorator(self):
        """Test detection when decorator name is not in import_aliases"""
        detector = JitDecoratorDetector()
        # Don't add anything to import_aliases
        
        # Parse @unknown_decorator as a decorator
        code = "@unknown_decorator\ndef func(): pass"
        tree = ast.parse(code)
        decorator_node = tree.body[0].decorator_list[0]
        
        # Should not detect as JIT decorator
        codeflash_output = detector._is_jit_decorator(decorator_node)

    def test_module_import_without_specific_decorator(self):
        """Test when only module is imported but not used as decorator directly"""
        detector = JitDecoratorDetector()
        detector.import_aliases["numba"] = ("numba", None)
        
        # Parse @numba as a decorator (just module name, not a decorator)
        code = "@numba\ndef func(): pass"
        tree = ast.parse(code)
        decorator_node = tree.body[0].decorator_list[0]
        
        # Should not detect as JIT decorator
        codeflash_output = detector._is_jit_decorator(decorator_node)

    def test_torch_compile_decorator(self):
        """Test detection of @torch.compile decorator"""
        detector = JitDecoratorDetector()
        detector.import_aliases["torch"] = ("torch", None)
        
        # Parse @torch.compile as a decorator
        code = "@torch.compile\ndef func(): pass"
        tree = ast.parse(code)
        decorator_node = tree.body[0].decorator_list[0]
        
        # Should detect torch.compile as a JIT decorator
        codeflash_output = detector._is_jit_decorator(decorator_node)

    def test_numba_cuda_jit_decorator(self):
        """Test detection of @numba.cuda.jit decorator"""
        detector = JitDecoratorDetector()
        detector.import_aliases["numba"] = ("numba", None)
        
        # Parse @numba.cuda.jit as a decorator
        code = "@numba.cuda.jit\ndef func(): pass"
        tree = ast.parse(code)
        decorator_node = tree.body[0].decorator_list[0]
        
        # Should detect numba.cuda.jit as a JIT decorator
        codeflash_output = detector._is_jit_decorator(decorator_node)

    def test_from_numba_import_jit(self):
        """Test detection when using: from numba import jit"""
        detector = JitDecoratorDetector()
        detector.import_aliases["jit"] = ("numba", "jit")
        
        # Parse @jit as a decorator
        code = "@jit\ndef func(): pass"
        tree = ast.parse(code)
        decorator_node = tree.body[0].decorator_list[0]
        
        # Should detect jit as a JIT decorator
        codeflash_output = detector._is_jit_decorator(decorator_node)

    def test_from_torch_jit_import_script(self):
        """Test detection when using: from torch.jit import script"""
        detector = JitDecoratorDetector()
        detector.import_aliases["script"] = ("torch.jit", "script")
        
        # Parse @script as a decorator
        code = "@script\ndef func(): pass"
        tree = ast.parse(code)
        decorator_node = tree.body[0].decorator_list[0]
        
        # Should detect script as a JIT decorator
        codeflash_output = detector._is_jit_decorator(decorator_node)

    
def func(): pass
"""
        tree = ast.parse(code)
        decorators = tree.body[0].decorator_list
        
        # Should detect both decorators
        codeflash_output = detector._is_jit_decorator(decorators[0])  # lru_cache
        codeflash_output = detector._is_jit_decorator(decorators[1])   # jit

    def test_check_name_decorator_batch(self):
        """Test _check_name_decorator with many names"""
        detector = JitDecoratorDetector()
        
        # Add 100 JIT imports and 100 non-JIT imports
        for i in range(100):
            detector.import_aliases[f"jit_{i}"] = ("numba", "jit")
            detector.import_aliases[f"other_{i}"] = ("functools", "other")
        
        # Test all JIT imports
        for i in range(100):
            pass
        
        # Test all non-JIT imports
        for i in range(100):
            pass

    def test_is_known_jit_decorator_batch(self):
        """Test _is_known_jit_decorator with many combinations"""
        detector = JitDecoratorDetector()
        
        # Test all combinations in JIT_DECORATORS
        jit_modules = list(detector._is_known_jit_decorator.__globals__["JIT_DECORATORS"].keys())

    def test_attribute_decorator_with_many_aliases(self):
        """Test attribute decorator detection with many import aliases"""
        detector = JitDecoratorDetector()
        
        # Add various module imports
        detector.import_aliases["nb"] = ("numba", None)
        detector.import_aliases["torch"] = ("torch", None)
        detector.import_aliases["jax"] = ("jax", None)
        detector.import_aliases["tf"] = ("tensorflow", None)
        
        # Test each one
        test_cases = [
            ("@nb.jit", True),
            ("@torch.compile", True),
            ("@torch.jit.script", True),
            ("@jax.jit", True),
            ("@tf.function", True),
        ]
        
        for decorator_str, expected in test_cases:
            code = f"{decorator_str}\ndef func(): pass"
            tree = ast.parse(code)
            decorator_node = tree.body[0].decorator_list[0]
            codeflash_output = detector._is_jit_decorator(decorator_node); result = codeflash_output

    def test_call_decorators_with_various_arguments(self):
        """Test call decorators with many different argument patterns"""
        detector = JitDecoratorDetector()
        detector.import_aliases["jit"] = ("numba", "jit")
        
        # Test various call patterns
        call_patterns = [
            "@jit()",
            "@jit(nopython=True)",
            "@jit(parallel=True)",
            "@jit('float64(float64)')",
            "@jit(fastmath=True, cache=True)",
            "@jit(target='cpu')",
        ]
        
        for decorator_str in call_patterns:
            code = f"{decorator_str}\ndef func(): pass"
            tree = ast.parse(code)
            decorator_node = tree.body[0].decorator_list[0]
            # All should be detected as JIT decorators
            codeflash_output = detector._is_jit_decorator(decorator_node)
# 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-pr1048-2026-01-14T01.50.34 and push.