⚡️ Speed up function _gridmake2_torch by 28%

[codeflash-ai]

📄 28% (0.28x) speedup for _gridmake2_torch in code_to_optimize/discrete_riccati.py

⏱️ Runtime : 987 microseconds → 774 microseconds (best of 72 runs)

📝 Explanation and details

The optimized code achieves a 27% speedup by eliminating expensive intermediate tensor allocations that were core to the original implementation.

Key Optimizations

1. Pre-allocation Instead of Concatenation

• Original: Created separate tensors via tile() and repeat_interleave(), then combined with column_stack() (3 allocations)
• Optimized: Pre-allocates the final output tensor once with torch.empty(), then fills it in-place using views
• Impact: Line profiler shows the original's column_stack took 20.4% (1D case) and 13.8% (2D case) of total time—completely eliminated

2. View-Based Assignment with Expand

• Original: x1.tile(n2) physically copies x1's data n2 times; x2.repeat_interleave(n1) physically replicates x2's elements
• Optimized: Uses unsqueeze().expand() which creates lightweight views (no data copy), then assigns through a reshaped view of the output buffer
• Why it's faster: expand() creates a view with modified strides—no memory copies until the final assignment, which writes directly to the pre-allocated output

3. Conditional Type Casting

• Only converts dtypes when necessary (x1c = x1 if x1.dtype == out_dtype else x1.to(out_dtype)), avoiding redundant conversions when types already match

Performance Analysis from Tests

The optimization particularly excels for larger inputs:

• test_large_scale_shapes_and_content_consistency (300×200): 38.6% faster (81.2μs → 58.6μs)
• test_large_scale_2d_x1_and_x2_1d_appending_behavior (100×5×50): 48.9% faster (28.9μs → 19.4μs)

For small inputs, the overhead of the additional setup (dtype promotion, device checks) slightly reduces gains or causes minor regressions in trivial cases (e.g., empty tensors: 25.6% slower). However, the overall benchmark shows a net 27% improvement across the mixed workload.

Behavior Preservation

The optimization maintains identical semantics:

• Same dtype promotion via torch.promote_types()
• Same exception raising for unsupported cases
• Identical output shapes and values (view-based filling produces the same cartesian product)
• No change to function signature or external behavior

✅ Correctness verification report:

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

⚙️ Click to see Existing Unit Tests

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_gridmake2_torch.py::TestGridmake2TorchCPU.test_2d_and_1d_simple	13.8μs	12.5μs	10.7%✅
test_gridmake2_torch.py::TestGridmake2TorchCPU.test_2d_and_1d_single_column	12.8μs	13.1μs	-2.22%⚠️
test_gridmake2_torch.py::TestGridmake2TorchCPU.test_both_1d_float_tensors	11.8μs	11.2μs	5.20%✅
test_gridmake2_torch.py::TestGridmake2TorchCPU.test_both_1d_simple	14.0μs	12.7μs	10.5%✅
test_gridmake2_torch.py::TestGridmake2TorchCPU.test_both_1d_single_element	12.5μs	11.2μs	11.9%✅
test_gridmake2_torch.py::TestGridmake2TorchCPU.test_large_tensors	19.5μs	15.8μs	23.8%✅
test_gridmake2_torch.py::TestGridmake2TorchCPU.test_not_implemented_for_1d_2d	1.21μs	1.29μs	-6.50%⚠️
test_gridmake2_torch.py::TestGridmake2TorchCPU.test_not_implemented_for_2d_2d	1.29μs	1.25μs	3.36%✅
test_gridmake2_torch.py::TestGridmake2TorchCPU.test_output_shape_1d_1d	12.2μs	12.2μs	0.345%✅
test_gridmake2_torch.py::TestGridmake2TorchCPU.test_output_shape_2d_1d	13.0μs	12.0μs	8.36%✅
test_gridmake2_torch.py::TestGridmake2TorchCPU.test_preserves_dtype_float64	12.3μs	11.5μs	7.64%✅
test_gridmake2_torch.py::TestGridmake2TorchCPU.test_preserves_dtype_int	14.2μs	12.4μs	14.5%✅

🌀 Click to see Generated Regression Tests

import pytest  # used for our unit tests
import torch  # torch is required for tensor operations

from code_to_optimize.discrete_riccati import _gridmake2_torch

# ----------------------------
# Test suite for _gridmake2_torch
# ----------------------------

# Basic tests (easiest -> harder)


def test_basic_1d_1d_integer_values():
    # Basic 1-D vs 1-D case with small integer tensors.
    x1 = torch.tensor([1, 2], dtype=torch.int64)  # 1-D
    x2 = torch.tensor([10, 20, 30], dtype=torch.int64)  # 1-D

    # Expected Cartesian product ordering based on implementation:
    # first column: x1 repeated len(x2) times -> [1,2,1,2,1,2]
    # second column: each x2 element repeated len(x1) times -> [10,10,20,20,30,30]
    expected = torch.tensor([[1, 10], [2, 10], [1, 20], [2, 20], [1, 30], [2, 30]], dtype=torch.int64)

    codeflash_output = _gridmake2_torch(x1, x2)
    out = codeflash_output  # 15.1μs -> 13.2μs (14.2% faster)


def test_basic_1d_1d_float_values_and_dtype_preservation():
    # 1-D float inputs: ensure float arithmetic and ordering are correct.
    x1 = torch.tensor([0.5, -1.0, 2.25], dtype=torch.float32)
    x2 = torch.tensor([100.0, 200.0], dtype=torch.float32)

    codeflash_output = _gridmake2_torch(x1, x2)
    out = codeflash_output  # 12.2μs -> 12.9μs (5.47% slower)


def test_basic_2d_1d_row_matrix_expansion():
    # x1 is 2-D (rows x columns), x2 is 1-D.
    # Check that rows of x1 get repeated and x2 becomes appended as an extra column.
    x1 = torch.tensor([[1.0, 1.5], [2.0, 2.5]], dtype=torch.float32)  # shape (2,2)
    x2 = torch.tensor([10.0, 20.0, 30.0], dtype=torch.float32)  # shape (3,)

    codeflash_output = _gridmake2_torch(x1, x2)
    out = codeflash_output  # 14.2μs -> 13.3μs (6.55% faster)


# Edge cases


def test_edge_both_1d_with_empty_input_one_empty():
    # Edge where one input is empty (length 0). Should return an empty tensor with 2 columns.
    x1 = torch.tensor([], dtype=torch.float32)  # length 0
    x2 = torch.tensor([1.0, 2.0], dtype=torch.float32)

    codeflash_output = _gridmake2_torch(x1, x2)
    out = codeflash_output  # 13.0μs -> 12.3μs (5.77% faster)


def test_edge_both_1d_with_both_empty():
    # Both inputs empty: expect empty (0 rows, 2 columns)
    x1 = torch.tensor([], dtype=torch.float64)
    x2 = torch.tensor([], dtype=torch.float64)

    codeflash_output = _gridmake2_torch(x1, x2)
    out = codeflash_output  # 8.46μs -> 11.4μs (25.6% slower)


def test_edge_2d_with_zero_rows_in_x1():
    # If x1 has zero rows but is 2-D (shape (0, ncols)) and x2 non-empty,
    # result should still be empty with columns = ncols + 1.
    x1 = torch.empty((0, 3), dtype=torch.float32)  # 0 rows, 3 columns
    x2 = torch.tensor([7.0, 8.0], dtype=torch.float32)  # 2 elements

    codeflash_output = _gridmake2_torch(x1, x2)
    out = codeflash_output  # 13.6μs -> 11.5μs (18.1% faster)


def test_raises_notimplemented_for_unsupported_combinations():
    # Cases where x2 is not 1-D, or x1 is 1-D while x2 is >1-D, should raise.
    x1 = torch.tensor([[1, 2], [3, 4]])  # 2-D
    x2 = torch.tensor([[10, 20], [30, 40]])  # 2-D -> unsupported

    with pytest.raises(NotImplementedError):
        _gridmake2_torch(x1, x2)  # 1.29μs -> 1.33μs (3.15% slower)

    # Also test 1-D x1 with 2-D x2 -> unsupported
    x1_1d = torch.tensor([1, 2, 3])
    x2_2d = torch.tensor([[1.0], [2.0]])
    with pytest.raises(NotImplementedError):
        _gridmake2_torch(x1_1d, x2_2d)  # 958ns -> 1.00μs (4.20% slower)


# Large scale tests (but stay under memory limits and avoid huge loops)
def test_large_scale_shapes_and_content_consistency():
    # Create reasonably large 1-D inputs but under memory constraints.
    # Using lengths 300 and 200 produces 60_000 rows; that's safe and well under 100MB.
    torch.manual_seed(0)  # deterministic randomness
    x1 = torch.arange(300, dtype=torch.float32)  # 0..299
    x2 = torch.arange(200, dtype=torch.float32) + 1000.0  # 1000..1199

    codeflash_output = _gridmake2_torch(x1, x2)
    out = codeflash_output  # 81.2μs -> 58.6μs (38.6% faster)

    # Check that the number of unique values in the first column equals len(x1)
    unique_first = torch.unique(out[:, 0])

    # Check that the set of unique first-column values equals x1's values
    # Sorting both for deterministic comparison
    sorted_unique_first, _ = torch.sort(unique_first)
    sorted_x1, _ = torch.sort(x1)

    # The element at row index (len(x1)-1) should be (x1[-1], x2[0]) due to ordering
    last_of_first_block_index = x1.numel() - 1


def test_large_scale_2d_x1_and_x2_1d_appending_behavior():
    # Larger 2-D x1 with many columns to ensure column appending works at scale.
    # Keep sizes reasonable: 100 rows x 5 cols repeated by x2 of length 50 -> 5000 rows.
    rows = 100
    cols = 5
    repeat_len = 50

    # Create x1 with identifiable pattern per row for verification
    # Each row i is filled with i (float), shape (rows, cols)
    x1 = torch.arange(rows, dtype=torch.float32).unsqueeze(1).repeat(1, cols)  # shape (rows, cols)
    x2 = torch.arange(repeat_len, dtype=torch.float32) * 10.0  # shape (repeat_len,)

    codeflash_output = _gridmake2_torch(x1, x2)
    out = codeflash_output  # 28.9μs -> 19.4μs (48.9% faster)


# Additional sanity: mixing dtypes but still 1-D x1 and x2 of same dtype
def test_mixed_sign_and_precision_behavior_consistency():
    # Verify function preserves order and arithmetic even when values include negatives and floats
    x1 = torch.tensor([-1.5, 0.0, 3.25], dtype=torch.float64)
    x2 = torch.tensor([2.5, -7.75], dtype=torch.float64)

    codeflash_output = _gridmake2_torch(x1, x2)
    out = codeflash_output  # 12.1μs -> 13.0μs (7.35% slower)

    # Validate that second column contains repeated-interleaved x2 entries in the expected order
    expected_second = x2.repeat_interleave(x1.numel())

    # And first column equals x1 repeated len(x2) times
    expected_first = x1.repeat(x2.numel())


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

import pytest
import torch

from code_to_optimize.discrete_riccati import _gridmake2_torch

# ============================================================================
# BASIC TEST CASES
# ============================================================================


class TestBasicFunctionality:
    """Test basic functionality of _gridmake2_torch with standard inputs."""

    def test_two_simple_1d_tensors(self):
        """Test cartesian product of two simple 1D tensors."""
        x1 = torch.tensor([1.0, 2.0])
        x2 = torch.tensor([3.0, 4.0])
        codeflash_output = _gridmake2_torch(x1, x2)
        result = codeflash_output  # 12.2μs -> 11.8μs (3.90% faster)

        # Verify values: x1 repeats as [1, 1, 2, 2] and x2 repeats as [3, 4, 3, 4]
        expected_first = torch.tensor([1.0, 1.0, 2.0, 2.0])
        expected_second = torch.tensor([3.0, 4.0, 3.0, 4.0])

    def test_1d_x1_larger_than_x2(self):
        """Test with x1 having more elements than x2."""
        x1 = torch.tensor([1.0, 2.0, 3.0])
        x2 = torch.tensor([10.0, 20.0])
        codeflash_output = _gridmake2_torch(x1, x2)
        result = codeflash_output  # 13.2μs -> 12.8μs (3.92% faster)

        expected_first = torch.tensor([1.0, 2.0, 3.0, 1.0, 2.0, 3.0])
        expected_second = torch.tensor([10.0, 10.0, 10.0, 20.0, 20.0, 20.0])

    def test_1d_x2_larger_than_x1(self):
        """Test with x2 having more elements than x1."""
        x1 = torch.tensor([5.0, 6.0])
        x2 = torch.tensor([100.0, 200.0, 300.0])
        codeflash_output = _gridmake2_torch(x1, x2)
        result = codeflash_output  # 12.5μs -> 12.4μs (0.330% faster)

        expected_first = torch.tensor([5.0, 6.0, 5.0, 6.0, 5.0, 6.0])
        expected_second = torch.tensor([100.0, 100.0, 200.0, 200.0, 300.0, 300.0])

    def test_2d_x1_with_1d_x2(self):
        """Test cartesian product of 2D tensor with 1D tensor."""
        x1 = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
        x2 = torch.tensor([10.0, 20.0])
        codeflash_output = _gridmake2_torch(x1, x2)
        result = codeflash_output  # 14.2μs -> 12.5μs (13.3% faster)

        # First two columns should be x1 repeated
        expected_x1_repeated = torch.tensor([[1.0, 2.0], [3.0, 4.0], [1.0, 2.0], [3.0, 4.0]])

        # Third column should be x2 repeated properly
        expected_x2 = torch.tensor([10.0, 10.0, 20.0, 20.0])

    def test_different_dtypes_float32(self):
        """Test with float32 tensors."""
        x1 = torch.tensor([1.0, 2.0], dtype=torch.float32)
        x2 = torch.tensor([3.0, 4.0], dtype=torch.float32)
        codeflash_output = _gridmake2_torch(x1, x2)
        result = codeflash_output  # 13.3μs -> 11.2μs (19.0% faster)

    def test_different_dtypes_float64(self):
        """Test with float64 tensors."""
        x1 = torch.tensor([1.0, 2.0], dtype=torch.float64)
        x2 = torch.tensor([3.0, 4.0], dtype=torch.float64)
        codeflash_output = _gridmake2_torch(x1, x2)
        result = codeflash_output  # 13.3μs -> 13.0μs (2.90% faster)


# ============================================================================
# EDGE CASE TEST CASES
# ============================================================================


class TestEdgeCases:
    """Test edge cases and boundary conditions."""

    def test_single_element_both_tensors(self):
        """Test with both tensors having a single element."""
        x1 = torch.tensor([5.0])
        x2 = torch.tensor([10.0])
        codeflash_output = _gridmake2_torch(x1, x2)
        result = codeflash_output  # 13.6μs -> 11.4μs (19.3% faster)

    def test_single_element_x1(self):
        """Test with x1 having a single element and x2 having multiple."""
        x1 = torch.tensor([7.0])
        x2 = torch.tensor([1.0, 2.0, 3.0])
        codeflash_output = _gridmake2_torch(x1, x2)
        result = codeflash_output  # 13.3μs -> 11.6μs (14.3% faster)
        expected_first = torch.tensor([7.0, 7.0, 7.0])
        expected_second = torch.tensor([1.0, 2.0, 3.0])

    def test_single_element_x2(self):
        """Test with x2 having a single element and x1 having multiple."""
        x1 = torch.tensor([1.0, 2.0, 3.0])
        x2 = torch.tensor([100.0])
        codeflash_output = _gridmake2_torch(x1, x2)
        result = codeflash_output  # 12.4μs -> 11.4μs (8.75% faster)
        expected_first = torch.tensor([1.0, 2.0, 3.0])
        expected_second = torch.tensor([100.0, 100.0, 100.0])

    def test_negative_values(self):
        """Test with negative values in tensors."""
        x1 = torch.tensor([-1.0, -2.0])
        x2 = torch.tensor([-10.0, -20.0])
        codeflash_output = _gridmake2_torch(x1, x2)
        result = codeflash_output  # 13.3μs -> 12.0μs (11.5% faster)
        expected_first = torch.tensor([-1.0, -1.0, -2.0, -2.0])
        expected_second = torch.tensor([-10.0, -20.0, -10.0, -20.0])

    def test_mixed_positive_negative(self):
        """Test with mixed positive and negative values."""
        x1 = torch.tensor([-1.0, 2.0, -3.0])
        x2 = torch.tensor([10.0, -5.0])
        codeflash_output = _gridmake2_torch(x1, x2)
        result = codeflash_output  # 13.4μs -> 12.6μs (6.29% faster)
        expected_first = torch.tensor([-1.0, 2.0, -3.0, -1.0, 2.0, -3.0])
        expected_second = torch.tensor([10.0, 10.0, 10.0, -5.0, -5.0, -5.0])

    def test_zero_values(self):
        """Test with zero values in tensors."""
        x1 = torch.tensor([0.0, 1.0, 2.0])
        x2 = torch.tensor([0.0, 5.0])
        codeflash_output = _gridmake2_torch(x1, x2)
        result = codeflash_output  # 11.7μs -> 11.4μs (2.57% faster)
        expected_first = torch.tensor([0.0, 1.0, 2.0, 0.0, 1.0, 2.0])
        expected_second = torch.tensor([0.0, 0.0, 0.0, 5.0, 5.0, 5.0])

    def test_very_small_values(self):
        """Test with very small floating point values."""
        x1 = torch.tensor([1e-6, 2e-6])
        x2 = torch.tensor([1e-7, 2e-7])
        codeflash_output = _gridmake2_torch(x1, x2)
        result = codeflash_output  # 11.6μs -> 11.1μs (4.89% faster)
        # Use looser tolerance for very small values
        expected_first = torch.tensor([1e-6, 1e-6, 2e-6, 2e-6])
        expected_second = torch.tensor([1e-7, 2e-7, 1e-7, 2e-7])

    def test_very_large_values(self):
        """Test with very large floating point values."""
        x1 = torch.tensor([1e6, 2e6])
        x2 = torch.tensor([1e7, 2e7])
        codeflash_output = _gridmake2_torch(x1, x2)
        result = codeflash_output  # 13.1μs -> 11.4μs (15.4% faster)
        expected_first = torch.tensor([1e6, 1e6, 2e6, 2e6])
        expected_second = torch.tensor([1e7, 2e7, 1e7, 2e7])

    def test_2d_x1_single_row_multiple_columns(self):
        """Test with 2D x1 having single row but multiple columns."""
        x1 = torch.tensor([[1.0, 2.0, 3.0]])
        x2 = torch.tensor([10.0, 20.0])
        codeflash_output = _gridmake2_torch(x1, x2)
        result = codeflash_output  # 14.5μs -> 13.0μs (10.9% faster)

        expected_x1_cols = torch.tensor([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]])

        expected_x2_col = torch.tensor([10.0, 20.0])

    def test_2d_x1_multiple_rows_single_column(self):
        """Test with 2D x1 having multiple rows but single column."""
        x1 = torch.tensor([[1.0], [2.0], [3.0]])
        x2 = torch.tensor([10.0, 20.0])
        codeflash_output = _gridmake2_torch(x1, x2)
        result = codeflash_output  # 14.2μs -> 12.3μs (15.6% faster)

        expected_first = torch.tensor([1.0, 2.0, 3.0, 1.0, 2.0, 3.0])
        expected_second = torch.tensor([10.0, 10.0, 10.0, 20.0, 20.0, 20.0])

    def test_2d_x1_with_3d_invalid_raises_error(self):
        """Test that 2D x1 with 2D x2 raises NotImplementedError."""
        x1 = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
        x2 = torch.tensor([[10.0, 20.0], [30.0, 40.0]])

        with pytest.raises(NotImplementedError):
            _gridmake2_torch(x1, x2)  # 1.33μs -> 1.29μs (3.25% faster)

    def test_1d_x1_with_2d_x2_raises_error(self):
        """Test that 1D x1 with 2D x2 raises NotImplementedError."""
        x1 = torch.tensor([1.0, 2.0])
        x2 = torch.tensor([[10.0, 20.0], [30.0, 40.0]])

        with pytest.raises(NotImplementedError):
            _gridmake2_torch(x1, x2)  # 1.21μs -> 1.21μs (0.000% faster)

To edit these changes git checkout codeflash/optimize-_gridmake2_torch-mkg28osu and push.