Fix strided_scan reading/writing out-of-bounds

[zcbenz]

The GPU strided_scan kernels can read/write out-of-bounds when the scanned axis is a slice. The test is from #4254.

Note that this is only a correctness fix, and does not try to optimize the mentioned case which currently has most of the threads wasted. At the moment I'm conservative on complicating the kernel for a crafted edge case.

[TheDarkchip]

Hi, the pointer forming here is UB (https://eel.is/c++draft/expr.add) ffor lanes where the result points beyond the position immediately after the buffer.

in += offset + global_index_x + read_offset_x;
out += offset + global_index_x + read_offset_x;

This is only pointer-formation UB, as the later guard correctly prevents the pointers from being dereferenced.

[zcbenz]

We don't quite care about undefined behaviors when writing GPU kernels, and UB is often abused when we can get a performance gain. Out-of-bound pointers are especially common in GPU kernels.

[TheDarkchip]

I see. I micro-benchmarked an exact-fit vector fast-path (only the Metal variant for now) TheDarkchip@58a2290 with the following script:

# bench_exact_fit.py
import argparse
import json
import statistics
import time

import mlx.core as mx


CASES = [
    ("matrix_128x256_axis0", (128, 256), 0),
    ("matrix_1024x256_axis0", (1024, 256), 0),
    ("matrix_4096x256_axis0", (4096, 256), 0),
    ("sequence_8x512x128_axis1", (8, 512, 128), 1),
    ("sequence_8x2048x256_axis1", (8, 2048, 256), 1),
    ("channels_32x32x1024_axis1", (32, 32, 1024), 1),
]


def execute(x, axis, inclusive):
    y = mx.cumsum(x, axis=axis, inclusive=inclusive)
    mx.eval(y)
    return y


def choose_iterations(x, axis, inclusive):
    for _ in range(8):
        execute(x, axis, inclusive)
    mx.synchronize()

    start = time.perf_counter_ns()
    for _ in range(10):
        execute(x, axis, inclusive)
    mx.synchronize()

    per_op = max((time.perf_counter_ns() - start) / 10, 1)
    target = round(80_000_000 / per_op)
    memory_cap = max(2, 32_000_000 // x.size)
    return max(2, min(2000, memory_cap, target))


def benchmark(name, shape, axis, inclusive):
    x = mx.ones(shape, dtype=mx.float32)
    mx.eval(x)
    mx.synchronize()
    iterations = choose_iterations(x, axis, inclusive)

    samples = []
    for _ in range(11):
        start = time.perf_counter_ns()
        for _ in range(iterations):
            y = execute(x, axis, inclusive)
        mx.synchronize()
        samples.append(
            (time.perf_counter_ns() - start) / iterations / 1000
        )

    axis_size = shape[axis]
    first = 1 if inclusive else 0
    expected_shape = [1] * len(shape)
    expected_shape[axis] = axis_size
    expected = mx.arange(
        first, first + axis_size, dtype=mx.float32
    ).reshape(expected_shape)

    if not bool(mx.all(y == expected).item()):
        raise RuntimeError(f"incorrect result for {name}")

    return {
        "name": name,
        "inclusive": inclusive,
        "median_us": statistics.median(samples),
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--variant", required=True)
    args = parser.parse_args()

    results = []
    for case in CASES:
        results.append(benchmark(*case, inclusive=True))
        results.append(benchmark(*case, inclusive=False))

    print(json.dumps({
        "variant": args.variant,
        "results": results,
    }))


if __name__ == "__main__":
    main()

Workload	Axis	Mode	#4430 (µs)	Minimal patch (µs)	Speedup
(128, 256)	0	Inclusive	253.1	248.9	1.02×
(128, 256)	0	Exclusive	248.2	259.3	0.96×
(1024, 256)	0	Inclusive	346.6	302.0	1.15×
(1024, 256)	0	Exclusive	342.7	310.2	1.10×
(4096, 256)	0	Inclusive	647.7	492.3	1.32×
(4096, 256)	0	Exclusive	668.9	511.7	1.31×
(8, 512, 128)	1	Inclusive	332.8	274.8	1.21×
(8, 512, 128)	1	Exclusive	330.0	284.6	1.16×
(8, 2048, 256)	1	Inclusive	687.2	577.2	1.19×
(8, 2048, 256)	1	Exclusive	807.6	607.8	1.33×
(32, 32, 1024)	1	Inclusive	458.1	435.8	1.05×
(32, 32, 1024)	1	Exclusive	478.2	463.3	1.03×

Overall geometric-mean speedup: 1.146×.
Geometric means from three alternating paired runs: 1.170×, 1.134×, and 1.159×.

May be worth looking into this in the future.