Faster cpu ops

[awni]

• Use a contiguous iterator for faster copy, unary, binary, ternary
• Some fixes for large arrays (use size_t in most places now)
• Remove some dim specializations to reduce bloat and binary size
• Add vector specialization for ternary
• Apply dimension collapsing for all the above ops

Benchmarks

Everything is way faster up to arbitrary dimensions. A few benchmarks:

Bench	Pre	Post
Copy 7D	5.811	5.867
Copy 8D	11.734	11.017
Binary	273.774	5.9
Unary	16.071	0.850

Note the copy is about the same because it had a 7D specialization. Now it only has up to 3D specialized.

mx.set_default_device(mx.cpu)

def copy(x):
    return x[..., ::-1]

def binary(x, y):
    return x + y

def unary(x):
    return mx.negative(x)

D = 10
x = mx.zeros((D, D, D, D, D, D, D))
x = mx.transpose(x, (1, 0, 4, 2, 3, 6, 5))
timeit(copy, x)

y = mx.zeros((D, D, D, D, D, D, D))
timeit(binary, x, y)

z = mx.random.uniform(shape=(12,)*6)[..., ::2]
mx.eval(unary(z))
timeit(unary, z)

[angeloskath]

It took some reading cause this moves some code that was added to fix an oob access in the copy. I am pretty sure it just moved further down in copy so no regression there let me know if I missed something.

Otherwise it looks great. Love to see all those red lines in the diff.

[awni]

It took some reading cause this moves some code that was added to fix an oob access in the copy.

Which code are you referring to? I can double check that I didn't mess that up.

[angeloskath]

This is the diff I was referring to. Git commit 03cf033 .

[angeloskath]

But it should be fine due to that.