Optimize waveform sampling with cross-platform SIMD

.github/workflows/workflow.yml

@@ -48,6 +48,10 @@ jobs:
     - name: Test with pytest
       run: |
         coverage run --source=waveforms -m pytest --verbose tests/
+    - name: Benchmark SIMD hot paths
+      if: ${{ matrix.python-version == '3.12' }}
+      run: |
+        python benchmarks/benchmark_simd.py --count 1000000 --repeats 7
     - name: Coveralls
       if: ${{ matrix.os == 'ubuntu-latest' && matrix.python-version == '3.12' }}
       continue-on-error: true

benchmarks/benchmark_simd.py

@@ -0,0 +1,80 @@
+"""Benchmark the loops targeted by native SIMD and batched-math kernels."""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+import time
+
+import numpy as np
+
+import waveforms as wf
+from waveforms._waveform import quantize_samples
+
+
+def best_time(function, repeats):
+    function()
+    timings = []
+    for _ in range(repeats):
+        gc.collect()
+        start = time.perf_counter()
+        function()
+        timings.append(time.perf_counter() - start)
+    return min(timings)
+
+
+def benchmark(count, rate, repeats):
+    duration = count / rate
+    start = -duration / 2
+    stop = start + duration
+    width = duration * 1.8
+    pulse = 0.8 * wf.gaussian(width) * wf.cos(2 * np.pi * 100e6)
+    powered = pulse**2
+    positions = np.linspace(start, stop, count, endpoint=False)
+    values = 0.999 * np.sin(np.linspace(-100, 100, count))
+
+    pulse.start = start
+    pulse.stop = stop
+    pulse.sample_rate = rate
+
+    cases = {
+        "evaluate_supported": lambda: pulse(positions),
+        "evaluate_power_fallback": lambda: powered(positions),
+        "sample_float64": lambda: pulse.sample(dtype=np.float64),
+        "sample_int16": lambda: pulse.sample(dtype=np.int16),
+        "sample_int32": lambda: pulse.sample(dtype=np.int32),
+        "quantize_int16": lambda: quantize_samples(values, 16),
+        "quantize_int32": lambda: quantize_samples(values, 32),
+    }
+    outputs = {name: function() for name, function in cases.items()}
+    if not np.array_equal(
+        outputs["sample_int16"], quantize_samples(outputs["sample_float64"], 16)
+    ):
+        raise AssertionError("int16 direct sampling differs from float quantization")
+    if not np.array_equal(
+        outputs["sample_int32"], quantize_samples(outputs["sample_float64"], 32)
+    ):
+        raise AssertionError("int32 direct sampling differs from float quantization")
+
+    return {
+        "count": count,
+        "rate": rate,
+        "milliseconds": {
+            name: 1e3 * best_time(function, repeats)
+            for name, function in cases.items()
+        },
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--count", type=int, default=1_000_000)
+    parser.add_argument("--rate", type=int, default=2_400_000_000)
+    parser.add_argument("--repeats", type=int, default=7)
+    args = parser.parse_args()
+    print(json.dumps(benchmark(args.count, args.rate, args.repeats), indent=2))
+
+
+if __name__ == "__main__":
+    main()

benchmarks/benchmark_simd_c.c

@@ -0,0 +1,154 @@
+#if !defined(_WIN32)
+#define _POSIX_C_SOURCE 200809L
+#endif
+
+#include "../waveforms/_cwaveform.h"
+
+#include <math.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#if defined(_WIN32)
+#define WIN32_LEAN_AND_MEAN
+#include <windows.h>
+#else
+#include <time.h>
+#endif
+
+static double now_seconds(void) {
+#if defined(_WIN32)
+    LARGE_INTEGER counter;
+    LARGE_INTEGER frequency;
+    QueryPerformanceFrequency(&frequency);
+    QueryPerformanceCounter(&counter);
+    return (double)counter.QuadPart / (double)frequency.QuadPart;
+#else
+    struct timespec value;
+    clock_gettime(CLOCK_MONOTONIC, &value);
+    return (double)value.tv_sec + 1e-9 * (double)value.tv_nsec;
+#endif
+}
+
+static double best_evaluate(const cwaveform_wave *wave, const double *positions,
+                            size_t count, double *output, int repeats) {
+    double best = INFINITY;
+    int repeat;
+    for (repeat = -1; repeat < repeats; ++repeat) {
+        double start = now_seconds();
+        double elapsed;
+        if (cwaveform_wave_evaluate(
+                wave, positions, count, 0, 1.0, -INFINITY, INFINITY,
+                output) != 0) abort();
+        elapsed = now_seconds() - start;
+        if (repeat >= 0 && elapsed < best) best = elapsed;
+    }
+    return best;
+}
+
+static double best_sample(const cwaveform_wave *wave, int64_t start_tick,
+                          size_t count, int dtype, void *output, int repeats) {
+    double best = INFINITY;
+    int repeat;
+    for (repeat = -1; repeat < repeats; ++repeat) {
+        double start = now_seconds();
+        double elapsed;
+        if (cwaveform_wave_sample(
+                wave, start_tick, count, 50, 1, 0, 1.0,
+                -INFINITY, INFINITY, dtype, 1.0, output) != 0) abort();
+        elapsed = now_seconds() - start;
+        if (repeat >= 0 && elapsed < best) best = elapsed;
+    }
+    return best;
+}
+
+static double best_quantize(const double *values, size_t count, int dtype,
+                            void *output, int repeats) {
+    double best = INFINITY;
+    int repeat;
+    for (repeat = -1; repeat < repeats; ++repeat) {
+        double start = now_seconds();
+        double elapsed;
+        if (cwaveform_quantize(values, count, dtype, 1.0, output) != 0) abort();
+        elapsed = now_seconds() - start;
+        if (repeat >= 0 && elapsed < best) best = elapsed;
+    }
+    return best;
+}
+
+int main(int argc, char **argv) {
+    const size_t count = argc > 1 ? (size_t)strtoull(argv[1], NULL, 10)
+                                  : 1000000u;
+    const int repeats = argc > 2 ? atoi(argv[2]) : 15;
+    const double rate = 2400000000.0;
+    const double duration = (double)count / rate;
+    const double start = -duration / 2.0;
+    const double pi = 3.14159265358979323846;
+    cwaveform_wave *gaussian = cwaveform_wave_gaussian(duration * 1.8);
+    cwaveform_wave *carrier = cwaveform_wave_cos(2.0 * pi * 100e6, 0.0);
+    cwaveform_wave *pulse = cwaveform_wave_mul_affine(
+        gaussian, 0, 0.8, carrier, 0, 1.0);
+    cwaveform_wave *powered = cwaveform_wave_power(pulse, 2);
+    double *positions = malloc(count * sizeof(*positions));
+    double *values = malloc(count * sizeof(*values));
+    double *samples = malloc(count * sizeof(*samples));
+    int16_t *samples16 = malloc(count * sizeof(*samples16));
+    int32_t *samples32 = malloc(count * sizeof(*samples32));
+    int16_t *expected16 = malloc(count * sizeof(*expected16));
+    int32_t *expected32 = malloc(count * sizeof(*expected32));
+    size_t index;
+    double evaluate_supported;
+    double evaluate_power;
+    double sample_float64;
+    double sample_int16;
+    double sample_int32;
+    double quantize_int16;
+    double quantize_int32;
+    if (pulse == NULL || powered == NULL || positions == NULL || values == NULL
+            || samples == NULL || samples16 == NULL || samples32 == NULL
+            || expected16 == NULL || expected32 == NULL) abort();
+    for (index = 0; index < count; ++index) {
+        positions[index] = start + (double)index / rate;
+        values[index] = 0.999 * sin(-100.0 + 200.0 * (double)index
+                                   / (double)count);
+    }
+    evaluate_supported = best_evaluate(
+        pulse, positions, count, samples, repeats);
+    evaluate_power = best_evaluate(powered, positions, count, samples, repeats);
+    sample_float64 = best_sample(
+        pulse, llround(start * 120000000000.0), count,
+        CWAVEFORM_FLOAT64, samples, repeats);
+    sample_int16 = best_sample(
+        pulse, llround(start * 120000000000.0), count,
+        CWAVEFORM_INT16, samples16, repeats);
+    sample_int32 = best_sample(
+        pulse, llround(start * 120000000000.0), count,
+        CWAVEFORM_INT32, samples32, repeats);
+    quantize_int16 = best_quantize(
+        values, count, CWAVEFORM_INT16, expected16, repeats);
+    quantize_int32 = best_quantize(
+        values, count, CWAVEFORM_INT32, expected32, repeats);
+    if (cwaveform_wave_sample(
+            pulse, llround(start * 120000000000.0), count, 50, 1, 0, 1.0,
+            -INFINITY, INFINITY, CWAVEFORM_FLOAT64, 1.0, samples) != 0
+            || cwaveform_quantize(
+                samples, count, CWAVEFORM_INT16, 1.0, expected16) != 0
+            || cwaveform_quantize(
+                samples, count, CWAVEFORM_INT32, 1.0, expected32) != 0
+            || memcmp(samples16, expected16, count * sizeof(*samples16)) != 0
+            || memcmp(samples32, expected32, count * sizeof(*samples32)) != 0)
+        abort();
+    printf("{\n  \"count\": %zu,\n  \"milliseconds\": {\n", count);
+    printf("    \"evaluate_supported\": %.6f,\n", 1e3 * evaluate_supported);
+    printf("    \"evaluate_power\": %.6f,\n", 1e3 * evaluate_power);
+    printf("    \"sample_float64\": %.6f,\n", 1e3 * sample_float64);
+    printf("    \"sample_int16\": %.6f,\n", 1e3 * sample_int16);
+    printf("    \"sample_int32\": %.6f,\n", 1e3 * sample_int32);
+    printf("    \"quantize_int16\": %.6f,\n", 1e3 * quantize_int16);
+    printf("    \"quantize_int32\": %.6f\n  }\n}\n", 1e3 * quantize_int32);
+    free(positions); free(values); free(samples); free(samples16);
+    free(samples32); free(expected16); free(expected32);
+    cwaveform_wave_release(powered); cwaveform_wave_release(pulse);
+    cwaveform_wave_release(carrier); cwaveform_wave_release(gaussian);
+    return 0;
+}

setup.py

@@ -39,15 +39,19 @@ def get_extensions():
                 sources = [os.path.join(dirpath, filename)]
                 include_dirs = []
                 extra_link_args = []
+                libraries = []
                 if filename == '_waveform.pyx':
                     sources.append(os.path.join(
                         'waveforms', '_cwaveform.c'))
                     include_dirs.append('waveforms')
                     if sys.platform == 'darwin':
                         extra_link_args.extend(['-framework', 'Accelerate'])
+                    elif sys.platform.startswith('linux'):
+                        libraries.append('m')
                 extensions.append(
                     Extension(module_name(dirpath, filename), sources,
                               include_dirs=include_dirs,
+                              libraries=libraries,
                               extra_link_args=extra_link_args))
 
     return extensions

tests/test_core.py

@@ -80,6 +80,78 @@ def test_c_symbolic_frequency_filter():
     )
 
 
+def test_c_symbolic_filter_reduces_trigonometric_products_and_powers():
+    x = np.linspace(-1.0, 1.0, 4097)
+    envelope = wf.gaussian(2.0)
+    first_frequency = 9.0
+    second_frequency = 7.0
+    first_phase = 0.31
+    second_phase = -0.27
+    cutoff = 5.0
+
+    first_cos = wf.cos(first_frequency, first_phase)
+    first_sin = wf.sin(first_frequency, first_phase)
+    second_cos = wf.cos(second_frequency, second_phase)
+    second_sin = wf.sin(second_frequency, second_phase)
+    difference = first_frequency - second_frequency
+    phase_difference = first_phase - second_phase
+
+    assert np.allclose(
+        (2 * envelope * first_cos * second_cos).filter(high=cutoff)(x),
+        (envelope * wf.cos(difference, phase_difference))(x),
+        atol=2e-11,
+    )
+    assert np.allclose(
+        (2 * envelope * first_cos * second_sin).filter(high=cutoff)(x),
+        (-envelope * wf.sin(difference, phase_difference))(x),
+        atol=2e-11,
+    )
+    assert np.allclose(
+        (2 * envelope * first_sin * second_sin).filter(high=cutoff)(x),
+        (envelope * wf.cos(difference, phase_difference))(x),
+        atol=2e-11,
+    )
+
+    powered = wf.cos(first_frequency, first_phase) ** 2
+    assert np.allclose(powered.filter(high=cutoff)(x), 0.5, atol=2e-11)
+    assert np.allclose(
+        powered.filter(2 * first_frequency, np.inf)(x),
+        (0.5 * wf.cos(2 * first_frequency, 2 * first_phase))(x),
+        atol=2e-11,
+    )
+
+    retained = (1.5e-15 * first_cos).filter(eps=1e-15)
+    discarded = (0.5e-15 * first_cos).filter(eps=1e-15)
+    assert np.max(np.abs(retained(x))) > 1e-15
+    assert np.array_equal(discarded(x), np.zeros_like(x))
+
+
+def test_down_conversion_filter_matches_historic_symbolic_behavior():
+    x = np.linspace(-100.0, 100.0, 10001)
+    envelope = wf.gaussian(100.0)
+    radio_frequency = 92.0451
+    local_frequency = 92.0
+    phase = 0.32
+    rf, _ = wf.mixing(
+        envelope, freq=radio_frequency, phase=phase, DRAGScaling=0.0,
+    )
+
+    i = (2 * rf * wf.cos(-2 * np.pi * local_frequency)).filter(
+        high=2 * np.pi * local_frequency,
+    )
+    q = (2 * rf * wf.sin(-2 * np.pi * local_frequency)).filter(
+        high=2 * np.pi * local_frequency,
+    )
+    difference = 2 * np.pi * (radio_frequency - local_frequency)
+
+    assert np.allclose(
+        i(x), (envelope * wf.cos(difference, -phase))(x), atol=2e-9,
+    )
+    assert np.allclose(
+        q(x), (envelope * wf.sin(difference, -phase))(x), atol=2e-9,
+    )
+
+
 def test_wave_block_roundtrip_pickle_and_numerics():
     actual = _pulse()
     actual.start = -20e-9

tests/test_waveform.py

@@ -473,6 +473,15 @@ def test_fixed_width_quantization_supported_sampling_and_fallback():
          1073741824, 2147483647, 2147483647],
     )
 
+    # Exercise the fused SIMD path, including its historic half-away-from-zero
+    # rounding rule and multidimensional ``out`` handling.
+    half_steps = np.array([0.5, -0.5, 1.5, -1.5] * 8) / 32768.0
+    target = np.empty((8, 4), dtype=np.int16)
+    assert quantize_samples(half_steps.reshape(8, 4), 16, out=target) is target
+    assert np.array_equal(
+        target.reshape(-1), np.array([1, -1, 2, -2] * 8, dtype=np.int16)
+    )
+
     wav = 0.8 * wf.gaussian(20e-9)
     wav.start = -20e-9
     wav.stop = 20e-9
@@ -493,6 +502,21 @@ def test_fixed_width_quantization_supported_sampling_and_fallback():
     assert np.array_equal(wav.sample(odd_rate), wav(legacy_grid))
 
 
+def test_simd_node_evaluation_matches_scalar_boundaries():
+    positions = np.linspace(-2.0, 2.0, 513)
+    waves = (
+        wf.t(),
+        wf.exp(0.2),
+        wf.sinc(1.3),
+        wf.cosh(0.3),
+        wf.sinh(0.3),
+        (wf.gaussian(4.0, 1.5) * wf.cos(2.2)) ** 2,
+    )
+    for wave in waves:
+        expected = np.array([wave(float(position)) for position in positions])
+        assert np.allclose(wave(positions), expected, rtol=1e-12, atol=1e-13)
+
+
 def test_integer_sampling_fast_paths_are_bit_exact_and_pickle_safe():
     rate = 2_400_000_000
     pulse = 0.8 * wf.gaussian(20e-9) * wf.cos(2 * np.pi * 100e6)