Faster cpu indexing

[awni]

Picking a little low hanging fruit.

Uses the ContiguousIterator in CPU scatter/gather for perf gains. Updated the iterator a little to make it easier to use.

Benchmarks

Timings on M1 Max. Source code below for benchmarks.

Bench	Pre	Post
Gather 1	228.555 (ms)	51.168 (ms)
Gather 2	139.326 (ms)	23.318 (ms)
Scatter 1	373.358 (ms)	128.734 (ms)
Scatter 2	321.287 (ms)	52.613 (ms)

D = 32
x = mx.zeros((D, D, D, D, D))
inds = [mx.zeros((D // 2, D // 2, D // 2), dtype=mx.int32).T] * 3

def fun(x, inds):
    return x[inds[0], :, inds[1], :, inds[2]]

timeit(fun, x, inds)

def fun(x, idx):
    return mx.take(x, idx, axis=0)

D = 32
x = mx.zeros((D, D, D, D, D))
x = x.transpose(1, 0, 4, 3, 2)
idx = mx.zeros((D // 2,)).astype(mx.int32)
timeit(fun, x, idx)

D = 32
x = mx.zeros((D, D, D, D, D))
inds = [mx.zeros((D // 2, D // 2, D // 2), dtype=mx.int32).T] * 3
update = mx.zeros((D // 2, D // 2, D, D, D // 2))
update = update.transpose(1, 0, 4, 3, 2)

def fun(x, inds, update):
    x[inds[0], :, inds[1], :, inds[2]] = update
    return x

timeit(fun, x, inds, update)

D = 32
x = mx.zeros((D, D, D, D, D))
x = x.transpose(1, 0, 4, 3, 2)
idx = mx.zeros((D // 2,)).astype(mx.int32)
update = mx.zeros((D, D // 2, D, D, D))
update = update.transpose(1, 0, 4, 3, 2)

def fun(x, idx, update):
    x[idx] = update
    return x

timeit(fun, x, idx, update)

[angeloskath]

Not much to say, it looks great! ContiguousIterator keeps on giving.