⚡️ Speed up function retrieve_batch_compatibility_of_input_selectors by 30% in PR #1504 (feature/try-to-beat-the-limitation-of-ee-in-terms-of-singular-elements-pushed-into-batch-inputs)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1504

If you approve this dependent PR, these changes will be merged into the original PR branch feature/try-to-beat-the-limitation-of-ee-in-terms-of-singular-elements-pushed-into-batch-inputs.

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

---

📄 30% (0.30x) speedup for retrieve_batch_compatibility_of_input_selectors in inference/core/workflows/execution_engine/v1/compiler/graph_constructor.py

⏱️ Runtime : 1.28 milliseconds → 987 microseconds (best of 274 runs)

📝 Explanation and details

The optimized code achieves a 29% speedup through two key optimizations that reduce overhead in the inner loop:

Key optimizations:

1. Eliminates repeated attribute lookups: Caches parsed_selector.definition.property_name in a local variable instead of accessing it twice per inner loop iteration
2. Reduces dictionary access overhead: Stores a reference to the target set (batch_compatibility_of_properties[property_name]) and reuses it, avoiding repeated dictionary lookups
3. Uses in-place set union (|=) instead of the update() method, which has slightly less overhead for set operations

Performance impact by test case:

• Small inputs (1-10 selectors): Modest 1-10% improvements due to reduced method call overhead
• Medium inputs (100-500 selectors): 12-25% speedups as the optimizations compound with more iterations
• Large inputs with many references: Up to 149% improvement in cases with many references per selector, where the inner loop dominates runtime

The line profiler shows the optimization moves expensive work (attribute lookups and dictionary access) from the inner loop to the outer loop. The original code performed parsed_selector.definition.property_name lookup 12,672 times, while the optimized version does it only 3,432 times - exactly once per selector instead of once per reference.

This optimization is particularly effective for workflows with selectors containing many allowed references, which is common in batch processing scenarios.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from collections import defaultdict
from typing import Dict, List, Set

# imports
import pytest  # used for our unit tests
from inference.core.workflows.execution_engine.v1.compiler.graph_constructor import \
    retrieve_batch_compatibility_of_input_selectors

# --- Minimal stubs for dependencies to allow testing ---

class Reference:
    def __init__(self, points_to_batch: Set[bool]):
        # points_to_batch is a set, e.g., {True}, {False}, or {True, False}
        self.points_to_batch = points_to_batch

class Definition:
    def __init__(self, property_name: str, allowed_references: List['Reference']):
        self.property_name = property_name
        self.allowed_references = allowed_references

class ParsedSelector:
    def __init__(self, definition: 'Definition'):
        self.definition = definition
from inference.core.workflows.execution_engine.v1.compiler.graph_constructor import \
    retrieve_batch_compatibility_of_input_selectors

# --- Unit Tests ---

# 1. Basic Test Cases

def test_single_selector_single_reference_true():
    # One selector, one reference, points_to_batch={True}
    ref = Reference({True})
    definition = Definition("foo", [ref])
    selector = ParsedSelector(definition)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([selector]); result = codeflash_output # 2.81μs -> 2.79μs (0.716% faster)

def test_single_selector_single_reference_false():
    # One selector, one reference, points_to_batch={False}
    ref = Reference({False})
    definition = Definition("bar", [ref])
    selector = ParsedSelector(definition)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([selector]); result = codeflash_output # 2.65μs -> 2.67μs (1.12% slower)

def test_single_selector_multiple_references_mixed():
    # One selector, multiple references, mixed True/False
    ref1 = Reference({True})
    ref2 = Reference({False})
    definition = Definition("baz", [ref1, ref2])
    selector = ParsedSelector(definition)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([selector]); result = codeflash_output # 2.85μs -> 2.73μs (4.02% faster)

def test_multiple_selectors_distinct_properties():
    # Two selectors, different property names
    ref1 = Reference({True})
    ref2 = Reference({False})
    def1 = Definition("alpha", [ref1])
    def2 = Definition("beta", [ref2])
    sel1 = ParsedSelector(def1)
    sel2 = ParsedSelector(def2)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel1, sel2]); result = codeflash_output # 3.15μs -> 3.12μs (1.25% faster)

def test_multiple_selectors_same_property():
    # Multiple selectors with the same property name
    ref1 = Reference({True})
    ref2 = Reference({False})
    def1 = Definition("gamma", [ref1])
    def2 = Definition("gamma", [ref2])
    sel1 = ParsedSelector(def1)
    sel2 = ParsedSelector(def2)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel1, sel2]); result = codeflash_output # 2.90μs -> 2.81μs (2.84% faster)

def test_selector_with_no_allowed_references():
    # Selector with no allowed references should yield empty set
    def1 = Definition("empty", [])
    sel1 = ParsedSelector(def1)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel1]); result = codeflash_output # 1.52μs -> 2.16μs (29.7% slower)

def test_empty_input():
    # No selectors at all
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([]); result = codeflash_output # 1.31μs -> 1.28μs (2.42% faster)

# 2. Edge Test Cases

def test_reference_with_empty_points_to_batch():
    # Reference with empty points_to_batch set
    ref = Reference(set())
    def1 = Definition("edge", [ref])
    sel = ParsedSelector(def1)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel]); result = codeflash_output # 2.48μs -> 2.44μs (1.68% faster)

def test_multiple_references_all_empty_points_to_batch():
    # All references have empty points_to_batch
    refs = [Reference(set()), Reference(set())]
    def1 = Definition("edge2", refs)
    sel = ParsedSelector(def1)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel]); result = codeflash_output # 2.56μs -> 2.50μs (2.40% faster)

def test_property_names_with_special_characters():
    # Property names with spaces and unicode
    ref = Reference({True})
    def1 = Definition("spaced name", [ref])
    def2 = Definition("unicodé", [ref])
    sel1 = ParsedSelector(def1)
    sel2 = ParsedSelector(def2)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel1, sel2]); result = codeflash_output # 3.15μs -> 3.14μs (0.319% faster)

def test_duplicate_references():
    # Duplicated references in allowed_references
    ref = Reference({True})
    def1 = Definition("dup", [ref, ref])
    sel = ParsedSelector(def1)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel]); result = codeflash_output # 2.79μs -> 2.65μs (5.29% faster)

def test_selector_with_mixed_empty_and_nonempty_references():
    # Some references empty, some not
    ref1 = Reference(set())
    ref2 = Reference({False})
    def1 = Definition("mix", [ref1, ref2])
    sel = ParsedSelector(def1)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel]); result = codeflash_output # 2.52μs -> 2.50μs (0.399% faster)

def test_property_name_collision_with_different_cases():
    # Property names that differ only in case
    ref1 = Reference({True})
    ref2 = Reference({False})
    def1 = Definition("Case", [ref1])
    def2 = Definition("case", [ref2])
    sel1 = ParsedSelector(def1)
    sel2 = ParsedSelector(def2)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel1, sel2]); result = codeflash_output # 2.96μs -> 2.88μs (2.77% faster)

# 3. Large Scale Test Cases

def test_large_number_of_selectors_unique_properties():
    # 500 selectors with unique property names, each with one reference (alternating True/False)
    selectors = []
    for i in range(500):
        ref = Reference({bool(i % 2)})
        defn = Definition(f"prop_{i}", [ref])
        selectors.append(ParsedSelector(defn))
    codeflash_output = retrieve_batch_compatibility_of_input_selectors(selectors); result = codeflash_output # 161μs -> 144μs (12.3% faster)
    # Each property should have a single value, alternating True/False
    for i in range(500):
        expected = {bool(i % 2)}

def test_large_number_of_selectors_shared_property():
    # 1000 selectors, all with the same property name, references alternate True/False
    selectors = []
    for i in range(1000):
        ref = Reference({bool(i % 2)})
        defn = Definition("bigprop", [ref])
        selectors.append(ParsedSelector(defn))
    codeflash_output = retrieve_batch_compatibility_of_input_selectors(selectors); result = codeflash_output # 162μs -> 129μs (25.5% faster)

def test_large_number_of_references_per_selector():
    # One selector, 1000 references, half True, half False
    refs = []
    for i in range(1000):
        refs.append(Reference({i % 2 == 0}))
    defn = Definition("huge", refs)
    sel = ParsedSelector(defn)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel]); result = codeflash_output # 90.3μs -> 36.2μs (149% faster)

def test_large_number_of_selectors_and_references_mixed():
    # 100 selectors, each with 10 references, mix of True/False
    selectors = []
    for i in range(100):
        refs = []
        for j in range(10):
            refs.append(Reference({(i + j) % 2 == 0}))
        defn = Definition(f"mix_{i}", refs)
        selectors.append(ParsedSelector(defn))
    codeflash_output = retrieve_batch_compatibility_of_input_selectors(selectors); result = codeflash_output # 126μs -> 64.8μs (95.0% faster)
    # Each property should have both True and False
    for i in range(100):
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from collections import defaultdict
from typing import Dict, List, Set

# imports
import pytest  # used for our unit tests
from inference.core.workflows.execution_engine.v1.compiler.graph_constructor import \
    retrieve_batch_compatibility_of_input_selectors


# ---- Mocked entities to simulate the real ones ----
# These are minimal mocks to allow the function and tests to run.
class MockReference:
    def __init__(self, points_to_batch):
        self.points_to_batch = points_to_batch  # Should be a set of bools

class MockDefinition:
    def __init__(self, property_name, allowed_references):
        self.property_name = property_name
        self.allowed_references = allowed_references  # List[MockReference]

class ParsedSelector:
    def __init__(self, definition):
        self.definition = definition
from inference.core.workflows.execution_engine.v1.compiler.graph_constructor import \
    retrieve_batch_compatibility_of_input_selectors

# ---- Unit Tests ----

# 1. Basic Test Cases

def test_empty_input_returns_empty_dict():
    # Test with no selectors: should return empty dict
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([]); result = codeflash_output # 1.32μs -> 1.37μs (3.64% slower)

def test_single_selector_single_reference_single_value():
    # One selector, one reference, one bool value
    ref = MockReference(points_to_batch={True})
    definition = MockDefinition(property_name="foo", allowed_references=[ref])
    selector = ParsedSelector(definition)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([selector]); result = codeflash_output # 2.65μs -> 2.56μs (3.12% faster)

def test_single_selector_single_reference_multiple_values():
    # One selector, one reference, multiple bool values
    ref = MockReference(points_to_batch={True, False})
    definition = MockDefinition(property_name="bar", allowed_references=[ref])
    selector = ParsedSelector(definition)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([selector]); result = codeflash_output # 2.47μs -> 2.37μs (4.61% faster)

def test_single_selector_multiple_references():
    # One selector, multiple references, overlapping values
    ref1 = MockReference(points_to_batch={True})
    ref2 = MockReference(points_to_batch={False})
    ref3 = MockReference(points_to_batch={True, False})
    definition = MockDefinition(property_name="baz", allowed_references=[ref1, ref2, ref3])
    selector = ParsedSelector(definition)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([selector]); result = codeflash_output # 3.03μs -> 2.75μs (10.2% faster)

def test_multiple_selectors_different_properties():
    # Multiple selectors, different property names
    ref1 = MockReference(points_to_batch={True})
    ref2 = MockReference(points_to_batch={False})
    def1 = MockDefinition(property_name="alpha", allowed_references=[ref1])
    def2 = MockDefinition(property_name="beta", allowed_references=[ref2])
    sel1 = ParsedSelector(def1)
    sel2 = ParsedSelector(def2)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel1, sel2]); result = codeflash_output # 3.02μs -> 3.05μs (1.02% slower)

def test_multiple_selectors_same_property():
    # Multiple selectors, same property name, should union values
    ref1 = MockReference(points_to_batch={True})
    ref2 = MockReference(points_to_batch={False})
    def1 = MockDefinition(property_name="gamma", allowed_references=[ref1])
    def2 = MockDefinition(property_name="gamma", allowed_references=[ref2])
    sel1 = ParsedSelector(def1)
    sel2 = ParsedSelector(def2)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel1, sel2]); result = codeflash_output # 2.88μs -> 2.88μs (0.000% faster)

# 2. Edge Test Cases

def test_selector_with_no_allowed_references():
    # Selector with no allowed references should result in empty set for property
    defn = MockDefinition(property_name="empty", allowed_references=[])
    sel = ParsedSelector(defn)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel]); result = codeflash_output # 1.50μs -> 2.11μs (28.9% slower)

def test_reference_with_empty_points_to_batch():
    # Reference with empty points_to_batch set
    ref = MockReference(points_to_batch=set())
    defn = MockDefinition(property_name="delta", allowed_references=[ref])
    sel = ParsedSelector(defn)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel]); result = codeflash_output # 2.36μs -> 2.33μs (1.29% faster)

def test_selector_with_duplicate_references():
    # Multiple references with the same points_to_batch
    ref = MockReference(points_to_batch={True})
    defn = MockDefinition(property_name="dup", allowed_references=[ref, ref, ref])
    sel = ParsedSelector(defn)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel]); result = codeflash_output # 2.94μs -> 2.65μs (10.5% faster)

def test_property_name_collision_with_different_cases():
    # Property names that differ only by case
    ref1 = MockReference(points_to_batch={True})
    ref2 = MockReference(points_to_batch={False})
    def1 = MockDefinition(property_name="Case", allowed_references=[ref1])
    def2 = MockDefinition(property_name="case", allowed_references=[ref2])
    sel1 = ParsedSelector(def1)
    sel2 = ParsedSelector(def2)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel1, sel2]); result = codeflash_output # 3.00μs -> 3.00μs (0.300% slower)

def test_selector_with_non_boolean_values_in_points_to_batch():
    # Reference with non-bool values (should not happen, but test for robustness)
    ref = MockReference(points_to_batch={1, 0, True, False})
    defn = MockDefinition(property_name="nonbool", allowed_references=[ref])
    sel = ParsedSelector(defn)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel]); result = codeflash_output # 2.34μs -> 2.33μs (0.817% faster)

def test_selector_with_none_in_points_to_batch():
    # Reference with None in points_to_batch
    ref = MockReference(points_to_batch={None})
    defn = MockDefinition(property_name="noneval", allowed_references=[ref])
    sel = ParsedSelector(defn)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel]); result = codeflash_output # 2.27μs -> 2.30μs (1.69% slower)

# 3. Large Scale Test Cases

def test_many_selectors_and_properties():
    # Test with 100 selectors, each with unique property and one reference
    selectors = []
    for i in range(100):
        ref = MockReference(points_to_batch={i % 2 == 0})
        defn = MockDefinition(property_name=f"prop_{i}", allowed_references=[ref])
        selectors.append(ParsedSelector(defn))
    codeflash_output = retrieve_batch_compatibility_of_input_selectors(selectors); result = codeflash_output # 38.6μs -> 35.6μs (8.42% faster)
    for i in range(100):
        expected = {True} if i % 2 == 0 else {False}

def test_many_references_per_selector():
    # Test with one selector, 100 references, alternating True/False
    refs = []
    for i in range(100):
        refs.append(MockReference(points_to_batch={i % 2 == 0}))
    defn = MockDefinition(property_name="bigprop", allowed_references=refs)
    sel = ParsedSelector(defn)
    codeflash_output = retrieve_batch_compatibility_of_input_selectors([sel]); result = codeflash_output # 11.3μs -> 6.13μs (84.8% faster)

def test_many_selectors_same_property_large_union():
    # 500 selectors, same property, each with different batch value
    selectors = []
    for i in range(500):
        val = bool(i % 2)
        ref = MockReference(points_to_batch={val})
        defn = MockDefinition(property_name="shared", allowed_references=[ref])
        selectors.append(ParsedSelector(defn))
    codeflash_output = retrieve_batch_compatibility_of_input_selectors(selectors); result = codeflash_output # 81.3μs -> 67.5μs (20.5% faster)

def test_large_mixture_of_empty_and_nonempty_references():
    # 200 selectors, half with empty references, half with non-empty
    selectors = []
    for i in range(200):
        if i % 2 == 0:
            defn = MockDefinition(property_name=f"p_{i}", allowed_references=[])
        else:
            ref = MockReference(points_to_batch={True})
            defn = MockDefinition(property_name=f"p_{i}", allowed_references=[ref])
        selectors.append(ParsedSelector(defn))
    codeflash_output = retrieve_batch_compatibility_of_input_selectors(selectors); result = codeflash_output # 40.3μs -> 55.2μs (27.1% slower)
    # Only odd-indexed properties should appear
    for i in range(200):
        pname = f"p_{i}"
        if i % 2 == 0:
            pass
        else:
            pass

def test_large_scale_performance():
    # Stress test: 1000 selectors, each with 2 references, each reference with {True, False}
    selectors = []
    for i in range(1000):
        refs = [MockReference(points_to_batch={True}), MockReference(points_to_batch={False})]
        defn = MockDefinition(property_name=f"prop_{i}", allowed_references=refs)
        selectors.append(ParsedSelector(defn))
    codeflash_output = retrieve_batch_compatibility_of_input_selectors(selectors); result = codeflash_output # 504μs -> 384μs (31.3% faster)
    for i in range(1000):
        pass
# 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-pr1504-2025-08-22T15.35.37 and push.