⚡️ Speed up function _prepare_for_blend by 63%

[codeflash-ai]

📄 63% (0.63x) speedup for _prepare_for_blend in src/diffusers/models/autoencoders/autoencoder_kl_allegro.py

⏱️ Runtime : 3.03 milliseconds → 1.85 milliseconds (best of 273 runs)

📝 Explanation and details

The main bottleneck in your code comes from repeatedly generating the 1D blend mask tensors with torch.arange(...).float().to(x.device) / overlap_x, followed by reshaping within every inner call. These are the lines where most time is spent and can be optimized.
Key idea: Precompute and cache the blend mask tensors for each overlap size seen during the runtime and reuse them.
We can add a helper to cache each blend mask tensor per (overlap, device) and direction ("start"/"end").

Summary of optimizations:

• Blend mask tensors are precomputed only once per overlap size, per device, per mask direction (start/reverse).
• No redundant arange computation and reshaping inside the main function--all done/cached in the helper.
• Minimizes device transfers and tensor allocation overhead.

This rewrite drastically reduces per-call runtime for the expensive masked multiplications.
The output is mathematically identical to your original code.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

import pytest  # used for our unit tests
import torch
from src.diffusers.models.autoencoders.autoencoder_kl_allegro import \
    _prepare_for_blend

# unit tests

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

def test_no_overlap_returns_input_unchanged():
    # All overlap params are 0, so tensor should remain unchanged
    x = torch.ones(1, 1, 4, 4, 4)
    x_orig = x.clone()
    codeflash_output = _prepare_for_blend((0, 1, 0), (0, 1, 0), (0, 1, 0), x); out = codeflash_output

def test_overlap_n_first_chunk():
    # Only n > 0 triggers first overlap_n region
    x = torch.ones(1, 1, 4, 4, 4)
    n_param = (1, 3, 2)  # n > 0, overlap_n=2
    codeflash_output = _prepare_for_blend(n_param, (0, 1, 0), (0, 1, 0), x.clone()); out = codeflash_output
    # First two slices along dim=2 should be scaled by [0/2, 1/2]
    expected = torch.ones(1, 1, 4, 4, 4)
    expected[:, :, 0, :, :] *= 0.0
    expected[:, :, 1, :, :] *= 0.5

def test_overlap_n_last_chunk():
    # Only n < n_max-1 triggers last overlap_n region
    x = torch.ones(1, 1, 4, 4, 4)
    n_param = (0, 3, 2)  # n < n_max-1, overlap_n=2
    codeflash_output = _prepare_for_blend(n_param, (0, 1, 0), (0, 1, 0), x.clone()); out = codeflash_output
    # Last two slices along dim=2 should be scaled by [1, 0.5]
    expected = torch.ones(1, 1, 4, 4, 4)
    expected[:, :, -2, :, :] *= 1.0
    expected[:, :, -1, :, :] *= 0.5

def test_overlap_h_middle_chunk():
    # Only h > 0 triggers first overlap_h region
    x = torch.ones(1, 1, 4, 4, 4)
    h_param = (1, 3, 2)
    codeflash_output = _prepare_for_blend((0, 1, 0), h_param, (0, 1, 0), x.clone()); out = codeflash_output
    # First two slices along dim=3 should be scaled by [0/2, 1/2]
    expected = torch.ones(1, 1, 4, 4, 4)
    expected[:, :, :, 0, :] *= 0.0
    expected[:, :, :, 1, :] *= 0.5

def test_overlap_h_last_chunk():
    # Only h < h_max-1 triggers last overlap_h region
    x = torch.ones(1, 1, 4, 4, 4)
    h_param = (0, 3, 2)
    codeflash_output = _prepare_for_blend((0, 1, 0), h_param, (0, 1, 0), x.clone()); out = codeflash_output
    # Last two slices along dim=3 should be scaled by [1, 0.5]
    expected = torch.ones(1, 1, 4, 4, 4)
    expected[:, :, :, -2, :] *= 1.0
    expected[:, :, :, -1, :] *= 0.5

def test_overlap_w_middle_chunk():
    # Only w > 0 triggers first overlap_w region
    x = torch.ones(1, 1, 4, 4, 4)
    w_param = (1, 3, 2)
    codeflash_output = _prepare_for_blend((0, 1, 0), (0, 1, 0), w_param, x.clone()); out = codeflash_output
    # First two slices along dim=4 should be scaled by [0/2, 1/2]
    expected = torch.ones(1, 1, 4, 4, 4)
    expected[:, :, :, :, 0] *= 0.0
    expected[:, :, :, :, 1] *= 0.5

def test_overlap_w_last_chunk():
    # Only w < w_max-1 triggers last overlap_w region
    x = torch.ones(1, 1, 4, 4, 4)
    w_param = (0, 3, 2)
    codeflash_output = _prepare_for_blend((0, 1, 0), (0, 1, 0), w_param, x.clone()); out = codeflash_output
    # Last two slices along dim=4 should be scaled by [1, 0.5]
    expected = torch.ones(1, 1, 4, 4, 4)
    expected[:, :, :, :, -2] *= 1.0
    expected[:, :, :, :, -1] *= 0.5

def test_all_overlaps_on_middle_chunk():
    # All overlaps active, and all > 0, so all "first" overlaps apply
    x = torch.ones(1, 1, 4, 4, 4)
    codeflash_output = _prepare_for_blend((1, 3, 2), (1, 3, 2), (1, 3, 2), x.clone()); out = codeflash_output
    # All first two slices along each axis should be scaled by [0, 0.5]
    expected = torch.ones(1, 1, 4, 4, 4)
    expected[:, :, 0, :, :] *= 0.0
    expected[:, :, 1, :, :] *= 0.5
    expected[:, :, :, 0, :] *= 0.0
    expected[:, :, :, 1, :] *= 0.5
    expected[:, :, :, :, 0] *= 0.0
    expected[:, :, :, :, 1] *= 0.5

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

def test_overlap_equals_tensor_size():
    # overlap equals the size of the dimension, so all elements are blended
    x = torch.ones(1, 1, 2, 2, 2)
    n_param = (1, 3, 2)
    codeflash_output = _prepare_for_blend(n_param, (0, 1, 0), (0, 1, 0), x.clone()); out = codeflash_output
    # All slices along dim=2 are blended: [0/2, 1/2]
    expected = torch.ones(1, 1, 2, 2, 2)
    expected[:, :, 0, :, :] *= 0.0
    expected[:, :, 1, :, :] *= 0.5

def test_overlap_larger_than_tensor_size():
    # If overlap is larger than dimension, should raise an error
    x = torch.ones(1, 1, 2, 2, 2)
    n_param = (1, 3, 3)  # overlap_n > size
    with pytest.raises(RuntimeError):
        _prepare_for_blend(n_param, (0, 1, 0), (0, 1, 0), x.clone())

def test_negative_overlap_raises():
    # Negative overlap should not be allowed
    x = torch.ones(1, 1, 4, 4, 4)
    n_param = (1, 3, -1)
    with pytest.raises(ValueError):
        # We'll add a check to the function for this test
        if n_param[2] < 0:
            raise ValueError("Negative overlap not allowed")
        _prepare_for_blend(n_param, (0, 1, 0), (0, 1, 0), x.clone())

def test_overlap_zero_does_nothing():
    # Overlap 0 should not change the tensor
    x = torch.rand(1, 1, 4, 4, 4)
    x_orig = x.clone()
    codeflash_output = _prepare_for_blend((0, 1, 0), (0, 1, 0), (0, 1, 0), x); out = codeflash_output

def test_tensor_with_different_dtype():
    # Should work with float32 and float64
    x32 = torch.ones(1, 1, 4, 4, 4, dtype=torch.float32)
    x64 = torch.ones(1, 1, 4, 4, 4, dtype=torch.float64)
    codeflash_output = _prepare_for_blend((1, 3, 2), (0, 1, 0), (0, 1, 0), x32.clone()); out32 = codeflash_output
    codeflash_output = _prepare_for_blend((1, 3, 2), (0, 1, 0), (0, 1, 0), x64.clone()); out64 = codeflash_output

def test_tensor_on_cuda_if_available():
    # Should work on CUDA tensors if available
    if torch.cuda.is_available():
        x = torch.ones(1, 1, 4, 4, 4).cuda()
        codeflash_output = _prepare_for_blend((1, 3, 2), (0, 1, 0), (0, 1, 0), x.clone()); out = codeflash_output
        # Check values as well
        expected = torch.ones(1, 1, 4, 4, 4).cuda()
        expected[:, :, 0, :, :] *= 0.0
        expected[:, :, 1, :, :] *= 0.5

def test_tensor_with_batch_and_channel_dims():
    # Test with batch size > 1 and channel > 1
    x = torch.ones(2, 3, 4, 4, 4)
    codeflash_output = _prepare_for_blend((1, 3, 2), (0, 1, 0), (0, 1, 0), x.clone()); out = codeflash_output
    expected = torch.ones(2, 3, 4, 4, 4)
    expected[:, :, 0, :, :] *= 0.0
    expected[:, :, 1, :, :] *= 0.5

def test_tensor_with_minimal_shape():
    # Minimal valid shape (1,1,1,1,1)
    x = torch.ones(1, 1, 1, 1, 1)
    codeflash_output = _prepare_for_blend((0, 1, 0), (0, 1, 0), (0, 1, 0), x.clone()); out = codeflash_output

def test_overlap_one():
    # overlap=1 is a special case: only first/last slice is affected
    x = torch.ones(1, 1, 4, 4, 4)
    n_param = (1, 3, 1)
    codeflash_output = _prepare_for_blend(n_param, (0, 1, 0), (0, 1, 0), x.clone()); out = codeflash_output
    expected = torch.ones(1, 1, 4, 4, 4)
    expected[:, :, 0, :, :] *= 0.0  # 0/1

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

def test_large_tensor_performance_and_correctness():
    # Large tensor, but < 100MB
    shape = (2, 2, 30, 30, 30)  # 2*2*30*30*30*4 bytes = ~432KB
    x = torch.ones(*shape)
    n_param = (1, 3, 10)
    h_param = (1, 3, 10)
    w_param = (1, 3, 10)
    codeflash_output = _prepare_for_blend(n_param, h_param, w_param, x.clone()); out = codeflash_output

def test_large_tensor_all_overlaps_zero():
    # Large tensor, all overlaps zero, should remain unchanged
    shape = (2, 2, 50, 50, 50)
    x = torch.rand(*shape)
    x_orig = x.clone()
    codeflash_output = _prepare_for_blend((0, 1, 0), (0, 1, 0), (0, 1, 0), x); out = codeflash_output

def test_large_tensor_overlap_equals_half_dim():
    # overlap equals half of the dimension size
    shape = (1, 1, 20, 20, 20)
    x = torch.ones(*shape)
    n_param = (1, 3, 10)
    codeflash_output = _prepare_for_blend(n_param, (0, 1, 0), (0, 1, 0), x.clone()); out = codeflash_output

def test_large_tensor_multiple_batches_channels():
    # Large tensor with multiple batches and channels
    shape = (5, 6, 10, 10, 10)
    x = torch.ones(*shape)
    n_param = (2, 3, 5)
    codeflash_output = _prepare_for_blend(n_param, (0, 1, 0), (0, 1, 0), x.clone()); out = codeflash_output
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest  # used for our unit tests
import torch
from src.diffusers.models.autoencoders.autoencoder_kl_allegro import \
    _prepare_for_blend

# unit tests

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

def test_no_overlaps_returns_input_unchanged():
    # Test when all overlaps are zero, tensor should remain unchanged
    x = torch.ones(1, 1, 2, 2, 2)
    x_clone = x.clone()
    codeflash_output = _prepare_for_blend((0, 1, 0), (0, 1, 0), (0, 1, 0), x); out = codeflash_output

def test_simple_overlap_n_start():
    # Test overlap in the n (depth) dimension at the start (n > 0)
    x = torch.ones(1, 1, 4, 2, 2)
    # overlap_n=2, n=1 (not first), n_max=3
    codeflash_output = _prepare_for_blend((1, 3, 2), (0, 1, 0), (0, 1, 0), x.clone()); out = codeflash_output
    # The first 2 slices along dim=2 should be scaled by [0/2, 1/2]
    expected = torch.ones(1, 1, 4, 2, 2)
    expected[:, :, 0, :, :] *= 0.0
    expected[:, :, 1, :, :] *= 0.5
    # The last 2 slices should be unchanged because n < n_max-1 triggers end overlap, but here n=1, n_max=3, so both start and end
    expected[:, :, -2, :, :] *= 0.5
    expected[:, :, -1, :, :] *= 0.0

def test_simple_overlap_n_end():
    # Test overlap in the n (depth) dimension at the end (n < n_max-1)
    x = torch.ones(1, 1, 4, 2, 2)
    # overlap_n=2, n=0 (first), n_max=3
    codeflash_output = _prepare_for_blend((0, 3, 2), (0, 1, 0), (0, 1, 0), x.clone()); out = codeflash_output
    expected = torch.ones(1, 1, 4, 2, 2)
    # Only end overlap should be applied
    expected[:, :, -2, :, :] *= 0.5
    expected[:, :, -1, :, :] *= 0.0

def test_simple_overlap_h_start():
    # Test overlap in the h (height) dimension at the start (h > 0)
    x = torch.ones(1, 1, 2, 4, 2)
    codeflash_output = _prepare_for_blend((0, 1, 0), (1, 3, 2), (0, 1, 0), x.clone()); out = codeflash_output
    expected = torch.ones(1, 1, 2, 4, 2)
    expected[:, :, :, 0, :] *= 0.0
    expected[:, :, :, 1, :] *= 0.5
    expected[:, :, :, -2, :] *= 0.5
    expected[:, :, :, -1, :] *= 0.0

def test_simple_overlap_w_start():
    # Test overlap in the w (width) dimension at the start (w > 0)
    x = torch.ones(1, 1, 2, 2, 4)
    codeflash_output = _prepare_for_blend((0, 1, 0), (0, 1, 0), (1, 3, 2), x.clone()); out = codeflash_output
    expected = torch.ones(1, 1, 2, 2, 4)
    expected[:, :, :, :, 0] *= 0.0
    expected[:, :, :, :, 1] *= 0.5
    expected[:, :, :, :, -2] *= 0.5
    expected[:, :, :, :, -1] *= 0.0

def test_overlap_all_dims():
    # Overlap in all dimensions, n>0, h>0, w>0
    x = torch.ones(1, 1, 4, 4, 4)
    codeflash_output = _prepare_for_blend((1, 3, 2), (1, 3, 2), (1, 3, 2), x.clone()); out = codeflash_output

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

def test_overlap_size_one():
    # Overlap size of 1 in each dimension (should scale to 0 at start/end)
    x = torch.ones(1, 1, 2, 2, 2)
    codeflash_output = _prepare_for_blend((1, 3, 1), (1, 3, 1), (1, 3, 1), x.clone()); out = codeflash_output

def test_overlap_size_equals_dim():
    # Overlap size equals dimension size (should blend entire dimension)
    x = torch.ones(1, 1, 4, 4, 4)
    codeflash_output = _prepare_for_blend((1, 3, 4), (1, 3, 4), (1, 3, 4), x.clone()); out = codeflash_output
    # All slices should be scaled by increasing/decreasing ramp from 0 to <1
    # For n-dim start
    expected_n = torch.arange(0, 4).float() / 4
    for i in range(4):
        pass
    # For h-dim start
    expected_h = torch.arange(0, 4).float() / 4
    for i in range(4):
        pass
    # For w-dim start
    expected_w = torch.arange(0, 4).float() / 4
    for i in range(4):
        pass

def test_zero_dim_size():
    # If any dimension is zero, should not crash, output shape should be preserved
    x = torch.ones(1, 1, 0, 4, 4)
    codeflash_output = _prepare_for_blend((0, 1, 0), (0, 1, 0), (0, 1, 0), x.clone()); out = codeflash_output
    x2 = torch.ones(1, 1, 4, 0, 4)
    codeflash_output = _prepare_for_blend((0, 1, 0), (0, 1, 0), (0, 1, 0), x2.clone()); out2 = codeflash_output
    x3 = torch.ones(1, 1, 4, 4, 0)
    codeflash_output = _prepare_for_blend((0, 1, 0), (0, 1, 0), (0, 1, 0), x3.clone()); out3 = codeflash_output


def test_negative_overlap():
    # Negative overlap should not apply any blending, should act as overlap=0
    x = torch.ones(1, 1, 2, 2, 2)
    codeflash_output = _prepare_for_blend((0, 1, -1), (0, 1, -1), (0, 1, -1), x.clone()); out = codeflash_output

def test_overlap_on_first_and_last_blocks():
    # Overlap only applies at start if n>0, at end if n<n_max-1
    x = torch.ones(1, 1, 4, 2, 2)
    # n=0 (first block), only end overlap should be applied
    codeflash_output = _prepare_for_blend((0, 3, 2), (0, 1, 0), (0, 1, 0), x.clone()); out = codeflash_output
    expected = torch.ones(1, 1, 4, 2, 2)
    expected[:, :, -2, :, :] *= 0.5
    expected[:, :, -1, :, :] *= 0.0
    # n=2 (last block), only start overlap should be applied
    codeflash_output = _prepare_for_blend((2, 3, 2), (0, 1, 0), (0, 1, 0), x.clone()); out2 = codeflash_output
    expected2 = torch.ones(1, 1, 4, 2, 2)
    expected2[:, :, 0, :, :] *= 0.0
    expected2[:, :, 1, :, :] *= 0.5

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

def test_large_tensor_single_overlap():
    # Large tensor, single overlap in n-dim
    x = torch.ones(2, 3, 100, 10, 10)
    codeflash_output = _prepare_for_blend((1, 3, 10), (0, 1, 0), (0, 1, 0), x.clone()); out = codeflash_output

def test_large_tensor_multiple_overlaps():
    # Large tensor, overlaps in all dims
    x = torch.ones(2, 3, 50, 50, 50)
    codeflash_output = _prepare_for_blend((1, 3, 10), (1, 3, 10), (1, 3, 10), x.clone()); out = codeflash_output

def test_large_tensor_edge_case_overlap_equals_dim():
    # Overlap equals dimension size for a large tensor
    x = torch.ones(1, 1, 100, 10, 10)
    codeflash_output = _prepare_for_blend((1, 3, 100), (0, 1, 0), (0, 1, 0), x.clone()); out = codeflash_output

def test_large_tensor_performance():
    # Large tensor but <100MB limit, check that function runs and output shape is preserved
    x = torch.ones(1, 1, 50, 50, 50)  # 1*1*50*50*50*4 bytes = 500KB
    codeflash_output = _prepare_for_blend((1, 3, 10), (1, 3, 10), (1, 3, 10), x.clone()); out = codeflash_output

def test_large_tensor_gradient():
    # Large tensor, check that blending ramps linearly in the overlap region
    x = torch.ones(1, 1, 20, 20, 20)
    codeflash_output = _prepare_for_blend((1, 3, 5), (0, 1, 0), (0, 1, 0), x.clone()); out = codeflash_output
    # n-dim start overlap: should be [0/5, 1/5, 2/5, 3/5, 4/5]
    for i in range(5):
        expected = i / 5.0
# 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-_prepare_for_blend-mbdjqwqv and push.