⚡️ Speed up method PositionEmbeddingSine.encode_points by 30%

[codeflash-ai]

📄 30% (0.30x) speedup for PositionEmbeddingSine.encode_points in ultralytics/models/sam/modules/blocks.py

⏱️ Runtime : 30.3 milliseconds → 23.3 milliseconds (best of 81 runs)

📝 Explanation and details

The optimization achieves a 29% speedup by replacing expensive tensor operations with more memory-efficient direct allocation and assignment patterns.

Key optimizations:

1. Eliminated torch.stack().flatten() operations: The original code used torch.stack((pos_x[:, 0::2].sin(), pos_x[:, 1::2].cos()), dim=2).flatten(1) which creates intermediate tensors and performs expensive reshape operations. The optimized version pre-allocates output tensors with torch.empty() and directly assigns values using slice operations (px[:, 0::2] = sin_x, px[:, 1::2] = cos_x).

2. Separated sin/cos computations: Instead of computing sin/cos within the stack operation, the optimized version computes them separately (sin_x = pos_x[:, 0::2].sin()), allowing for better memory locality and avoiding redundant computations.

3. Direct memory assignment: Using pre-allocated tensors and slice assignment is significantly faster than tensor concatenation and reshaping operations, especially for larger tensors.

Performance impact from test results:

• Small inputs show modest slowdowns (3-6% in basic tests) due to the overhead of separate allocations, but this is negligible in absolute terms (microseconds).
• Large-scale tests show substantial improvements: 77% faster for 100x100 batches, 33% faster for large embedding dimensions, and 32% faster for stress tests with maximum memory usage.

The optimization is particularly effective for the SAM (Segment Anything Model) use case where position embeddings are computed for large batches of image coordinates, making this a valuable improvement for computer vision workloads that process many spatial positions simultaneously.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

# imports
import pytest  # used for our unit tests
import torch
from ultralytics.models.sam.modules.blocks import PositionEmbeddingSine

# unit tests

# --------------------- Basic Test Cases ---------------------


def test_encode_points_basic_shape_and_dtype():
    # Basic test: check output shape and dtype for small input
    pe = PositionEmbeddingSine(num_pos_feats=8)
    x = torch.tensor([[0.1, 0.2], [0.3, 0.4]])
    y = torch.tensor([[0.5, 0.6], [0.7, 0.8]])
    labels = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
    codeflash_output = pe.encode_points(x, y, labels)
    out = codeflash_output  # 142μs -> 148μs (3.72% slower)


def test_encode_points_label_preservation():
    # Check that the label is preserved in the last column
    pe = PositionEmbeddingSine(num_pos_feats=8)
    x = torch.tensor([[0.0, 0.5]])
    y = torch.tensor([[0.0, 0.5]])
    labels = torch.tensor([[42.0, -1.0]])
    codeflash_output = pe.encode_points(x, y, labels)
    out = codeflash_output  # 135μs -> 136μs (0.607% slower)


def test_encode_points_batch_point_consistency():
    # Check that batches and points are handled correctly
    pe = PositionEmbeddingSine(num_pos_feats=16)
    batch = 3
    points = 5
    x = torch.arange(batch * points, dtype=torch.float32).reshape(batch, points) / 10
    y = torch.arange(batch * points, dtype=torch.float32).reshape(batch, points) / 20
    labels = torch.arange(batch * points, dtype=torch.float32).reshape(batch, points)
    codeflash_output = pe.encode_points(x, y, labels)
    out = codeflash_output  # 109μs -> 113μs (3.92% slower)


def test_encode_points_even_num_pos_feats():
    # Check that odd num_pos_feats raises AssertionError
    with pytest.raises(AssertionError):
        PositionEmbeddingSine(num_pos_feats=7)


def test_encode_points_value_range():
    # Check that output values are bounded between -1 and 1 except for label
    pe = PositionEmbeddingSine(num_pos_feats=8)
    x = torch.tensor([[0.0, 1.0]])
    y = torch.tensor([[0.0, 1.0]])
    labels = torch.tensor([[0.0, 1.0]])
    codeflash_output = pe.encode_points(x, y, labels)
    out = codeflash_output  # 142μs -> 151μs (5.97% slower)


# --------------------- Edge Test Cases ---------------------


def test_encode_points_empty_batch():
    # Edge: Empty batch
    pe = PositionEmbeddingSine(num_pos_feats=8)
    x = torch.empty((0, 3))
    y = torch.empty((0, 3))
    labels = torch.empty((0, 3))
    codeflash_output = pe.encode_points(x, y, labels)
    out = codeflash_output


def test_encode_points_empty_points():
    # Edge: Zero points per batch
    pe = PositionEmbeddingSine(num_pos_feats=8)
    x = torch.empty((2, 0))
    y = torch.empty((2, 0))
    labels = torch.empty((2, 0))
    codeflash_output = pe.encode_points(x, y, labels)
    out = codeflash_output


def test_encode_points_mismatched_shapes():
    # Edge: Mismatched batch/point shapes should raise AssertionError
    pe = PositionEmbeddingSine(num_pos_feats=8)
    x = torch.zeros((2, 3))
    y = torch.zeros((2, 4))
    labels = torch.zeros((2, 3))
    with pytest.raises(AssertionError):
        pe.encode_points(x, y, labels)  # 10.6μs -> 10.7μs (1.20% slower)


def test_encode_points_extreme_values():
    # Edge: Coordinates with extreme values
    pe = PositionEmbeddingSine(num_pos_feats=8)
    x = torch.tensor([[1e6, -1e6]])
    y = torch.tensor([[1e6, -1e6]])
    labels = torch.tensor([[0.0, 1.0]])
    codeflash_output = pe.encode_points(x, y, labels)
    out = codeflash_output  # 155μs -> 160μs (2.67% slower)


def test_encode_points_inf_nan_inputs():
    # Edge: Inputs containing inf/nan should propagate
    pe = PositionEmbeddingSine(num_pos_feats=8)
    x = torch.tensor([[float("inf"), float("nan")]])
    y = torch.tensor([[float("inf"), float("nan")]])
    labels = torch.tensor([[1.0, 2.0]])
    codeflash_output = pe.encode_points(x, y, labels)
    out = codeflash_output  # 137μs -> 146μs (6.25% slower)


def test_encode_points_scale_and_normalize_behavior():
    # Edge: Passing scale when normalize is False should raise ValueError
    with pytest.raises(ValueError):
        PositionEmbeddingSine(num_pos_feats=8, normalize=False, scale=3.14)


def test_encode_points_device_consistency():
    # Edge: Inputs on CUDA (if available) should work and output on same device
    if torch.cuda.is_available():
        pe = PositionEmbeddingSine(num_pos_feats=8).cuda()
        x = torch.tensor([[0.1, 0.2]], device="cuda")
        y = torch.tensor([[0.3, 0.4]], device="cuda")
        labels = torch.tensor([[1.0, 2.0]], device="cuda")
        codeflash_output = pe.encode_points(x, y, labels)
        out = codeflash_output


def test_encode_points_non_float_labels():
    # Edge: Labels are integer type, output should convert to float
    pe = PositionEmbeddingSine(num_pos_feats=8)
    x = torch.tensor([[0.1, 0.2]])
    y = torch.tensor([[0.3, 0.4]])
    labels = torch.tensor([[1, 2]], dtype=torch.int64)
    codeflash_output = pe.encode_points(x, y, labels)
    out = codeflash_output  # 164μs -> 168μs (2.29% slower)


# --------------------- Large Scale Test Cases ---------------------


def test_encode_points_large_batch_and_points():
    # Large scale: 100 batches, 100 points each
    pe = PositionEmbeddingSine(num_pos_feats=32)
    batch = 100
    points = 100
    x = torch.rand((batch, points), dtype=torch.float32)
    y = torch.rand((batch, points), dtype=torch.float32)
    labels = torch.arange(batch * points, dtype=torch.float32).reshape(batch, points)
    codeflash_output = pe.encode_points(x, y, labels)
    out = codeflash_output  # 1.21ms -> 682μs (77.2% faster)


def test_encode_points_large_embedding_dim():
    # Large scale: large embedding dimension (but < 100MB output)
    num_pos_feats = 128  # 128*2+1=257 floats per point
    batch = 5
    points = 100
    pe = PositionEmbeddingSine(num_pos_feats=num_pos_feats)
    x = torch.rand((batch, points), dtype=torch.float32)
    y = torch.rand((batch, points), dtype=torch.float32)
    labels = torch.zeros((batch, points), dtype=torch.float32)
    codeflash_output = pe.encode_points(x, y, labels)
    out = codeflash_output  # 319μs -> 239μs (33.3% faster)


def test_encode_points_stress_max_memory():
    # Large scale: maximize tensor size under 100MB
    # Each float32 is 4 bytes, so 100MB/4 = 25,000,000 floats
    # Let's use batch=10, points=1000, num_pos_feats=512 (so output shape [10,1000,1025], ~41MB)
    batch = 10
    points = 1000
    num_pos_feats = 512
    pe = PositionEmbeddingSine(num_pos_feats=num_pos_feats)
    x = torch.rand((batch, points), dtype=torch.float32)
    y = torch.rand((batch, points), dtype=torch.float32)
    labels = torch.zeros((batch, points), dtype=torch.float32)
    codeflash_output = pe.encode_points(x, y, labels)
    out = codeflash_output  # 25.7ms -> 19.3ms (32.8% faster)


def test_encode_points_grad_disabled():
    # Large scale: ensure @torch.no_grad disables grads
    pe = PositionEmbeddingSine(num_pos_feats=16)
    x = torch.rand((2, 2), dtype=torch.float32, requires_grad=True)
    y = torch.rand((2, 2), dtype=torch.float32, requires_grad=True)
    labels = torch.rand((2, 2), dtype=torch.float32, requires_grad=True)
    codeflash_output = pe.encode_points(x, y, labels)
    out = codeflash_output  # 145μs -> 149μs (2.43% slower)


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

# imports
import pytest  # used for our unit tests
import torch  # required for tensor operations
from ultralytics.models.sam.modules.blocks import PositionEmbeddingSine

# unit tests

# --------------- BASIC TEST CASES ---------------


def test_encode_points_basic_shape_and_type():
    # Test that output shape and type are correct for a simple input
    pos_emb = PositionEmbeddingSine(num_pos_feats=8)
    # 1 batch, 2 points
    x = torch.tensor([[0.1, 0.2]], dtype=torch.float32)
    y = torch.tensor([[0.3, 0.4]], dtype=torch.float32)
    labels = torch.tensor([[1, 2]], dtype=torch.float32)
    codeflash_output = pos_emb.encode_points(x, y, labels)
    out = codeflash_output  # 135μs -> 142μs (5.02% slower)


def test_encode_points_batch_multiple_points():
    # Test with batch size > 1 and multiple points
    pos_emb = PositionEmbeddingSine(num_pos_feats=4)
    x = torch.tensor([[0.0, 0.5], [1.0, 1.5]], dtype=torch.float32)
    y = torch.tensor([[0.1, 0.6], [1.1, 1.6]], dtype=torch.float32)
    labels = torch.tensor([[0, 1], [2, 3]], dtype=torch.float32)
    codeflash_output = pos_emb.encode_points(x, y, labels)
    out = codeflash_output  # 136μs -> 141μs (3.56% slower)


def test_encode_points_values_change_with_input():
    # The output should change if the input coordinates change
    pos_emb = PositionEmbeddingSine(num_pos_feats=4)
    x1 = torch.tensor([[0.1, 0.2]], dtype=torch.float32)
    y1 = torch.tensor([[0.3, 0.4]], dtype=torch.float32)
    labels = torch.tensor([[1, 2]], dtype=torch.float32)
    codeflash_output = pos_emb.encode_points(x1, y1, labels)
    out1 = codeflash_output  # 135μs -> 141μs (4.16% slower)
    x2 = torch.tensor([[0.5, 0.6]], dtype=torch.float32)
    y2 = torch.tensor([[0.7, 0.8]], dtype=torch.float32)
    codeflash_output = pos_emb.encode_points(x2, y2, labels)
    out2 = codeflash_output  # 68.6μs -> 75.0μs (8.59% slower)


def test_encode_points_label_in_last_column():
    # The last column should always be the label
    pos_emb = PositionEmbeddingSine(num_pos_feats=4)
    x = torch.tensor([[0.1, 0.2]], dtype=torch.float32)
    y = torch.tensor([[0.3, 0.4]], dtype=torch.float32)
    labels = torch.tensor([[5, 6]], dtype=torch.float32)
    codeflash_output = pos_emb.encode_points(x, y, labels)
    out = codeflash_output  # 131μs -> 138μs (5.31% slower)


# --------------- EDGE TEST CASES ---------------


def test_encode_points_minimal_input():
    # Minimal input: batch=1, n_points=1
    pos_emb = PositionEmbeddingSine(num_pos_feats=2)
    x = torch.tensor([[0.0]], dtype=torch.float32)
    y = torch.tensor([[0.0]], dtype=torch.float32)
    labels = torch.tensor([[0]], dtype=torch.float32)
    codeflash_output = pos_emb.encode_points(x, y, labels)
    out = codeflash_output


def test_encode_points_large_coordinates():
    # Test with very large coordinates
    pos_emb = PositionEmbeddingSine(num_pos_feats=4)
    x = torch.tensor([[1e6, -1e6]], dtype=torch.float32)
    y = torch.tensor([[1e7, -1e7]], dtype=torch.float32)
    labels = torch.tensor([[0, 1]], dtype=torch.float32)
    codeflash_output = pos_emb.encode_points(x, y, labels)
    out = codeflash_output  # 155μs -> 161μs (3.87% slower)


def test_encode_points_negative_and_zero_coordinates():
    # Test with negative and zero coordinates
    pos_emb = PositionEmbeddingSine(num_pos_feats=4)
    x = torch.tensor([[0.0, -0.5]], dtype=torch.float32)
    y = torch.tensor([[0.0, -1.0]], dtype=torch.float32)
    labels = torch.tensor([[1, 2]], dtype=torch.float32)
    codeflash_output = pos_emb.encode_points(x, y, labels)
    out = codeflash_output  # 135μs -> 144μs (6.21% slower)


def test_encode_points_non_float_labels():
    # Test with integer labels (should be castable to float)
    pos_emb = PositionEmbeddingSine(num_pos_feats=4)
    x = torch.tensor([[0.1, 0.2]], dtype=torch.float32)
    y = torch.tensor([[0.3, 0.4]], dtype=torch.float32)
    labels = torch.tensor([[1, 2]], dtype=torch.int32).float()
    codeflash_output = pos_emb.encode_points(x, y, labels)
    out = codeflash_output  # 127μs -> 130μs (2.61% slower)


def test_encode_points_mismatched_shapes_raises():
    # Should raise assertion if shapes do not match
    pos_emb = PositionEmbeddingSine(num_pos_feats=4)
    x = torch.tensor([[0.1, 0.2]], dtype=torch.float32)
    y = torch.tensor([[0.3, 0.4]], dtype=torch.float32)
    labels = torch.tensor([[1, 2, 3]], dtype=torch.float32)  # Wrong shape
    with pytest.raises(AssertionError):
        pos_emb.encode_points(x, y, labels)  # 9.62μs -> 9.26μs (3.89% faster)


def test_encode_points_odd_num_pos_feats_assertion():
    # Should raise assertion if num_pos_feats is not even
    with pytest.raises(AssertionError):
        PositionEmbeddingSine(num_pos_feats=7)


def test_encode_points_scale_requires_normalize():
    # Should raise ValueError if scale is set but normalize is False
    with pytest.raises(ValueError):
        PositionEmbeddingSine(num_pos_feats=8, scale=1.0, normalize=False)


# --------------- LARGE SCALE TEST CASES ---------------


def test_encode_points_large_batch_and_points():
    # Test with large batch and number of points, but <100MB
    pos_emb = PositionEmbeddingSine(num_pos_feats=32)
    batch_size = 16
    num_points = 32
    x = torch.rand((batch_size, num_points), dtype=torch.float32)
    y = torch.rand((batch_size, num_points), dtype=torch.float32)
    labels = torch.randint(0, 10, (batch_size, num_points), dtype=torch.float32)
    codeflash_output = pos_emb.encode_points(x, y, labels)
    out = codeflash_output  # 194μs -> 184μs (5.50% faster)


def test_encode_points_maximum_allowed_tensor_size():
    # Test with tensor size just below 100MB
    num_pos_feats = 64  # 64*2+1 = 129
    batch_size = 32
    num_points = 24  # 32*24*129*4 bytes = ~397kB
    pos_emb = PositionEmbeddingSine(num_pos_feats=num_pos_feats)
    x = torch.rand((batch_size, num_points), dtype=torch.float32)
    y = torch.rand((batch_size, num_points), dtype=torch.float32)
    labels = torch.randint(0, 100, (batch_size, num_points), dtype=torch.float32)
    codeflash_output = pos_emb.encode_points(x, y, labels)
    out = codeflash_output  # 269μs -> 220μs (22.3% faster)


def test_encode_points_large_scale_consistency():
    # Ensure that output is consistent for repeated calls with same input
    pos_emb = PositionEmbeddingSine(num_pos_feats=16)
    batch_size = 8
    num_points = 32
    x = torch.rand((batch_size, num_points), dtype=torch.float32)
    y = torch.rand((batch_size, num_points), dtype=torch.float32)
    labels = torch.randint(0, 10, (batch_size, num_points), dtype=torch.float32)
    codeflash_output = pos_emb.encode_points(x, y, labels)
    out1 = codeflash_output  # 156μs -> 156μs (0.139% slower)
    codeflash_output = pos_emb.encode_points(x, y, labels)
    out2 = codeflash_output  # 85.0μs -> 86.0μs (1.08% slower)


def test_encode_points_large_scale_label_column():
    # For large input, last column should still match labels
    pos_emb = PositionEmbeddingSine(num_pos_feats=32)
    batch_size = 16
    num_points = 32
    x = torch.rand((batch_size, num_points), dtype=torch.float32)
    y = torch.rand((batch_size, num_points), dtype=torch.float32)
    labels = torch.randint(0, 100, (batch_size, num_points), dtype=torch.float32)
    codeflash_output = pos_emb.encode_points(x, y, labels)
    out = codeflash_output  # 176μs -> 163μs (7.79% faster)
    # Check last column for every batch
    for b in range(batch_size):
        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-PositionEmbeddingSine.encode_points-mj9u30kz and push.