feat: simplify metal reduce kernels and standardize on u32 indexing

[drbh]

This PR improves the metal_reduce_.*_strided kernels

Benchmark

cargo bench --features metal -p candle-core --bench bench_main -- "metal_reduce_.*_strided"

Results

| Benchmark    | main                | new                      | Speedup      |
|--------------|---------------------|--------------------------|--------------|
| f32_strided  | 5.62ms (711 MiB/s)  | 2.3433 ms (1.6670 GiB/s) | 2.39x faster |
| f16_strided  | 3.98ms (503 MiB/s)  | 2.3631 ms (851.33 MiB/s) | 1.69x faster |
| bf16_strided | 3.95ms (505 MiB/s)  | 2.3545 ms (849.44 MiB/s) | 1.68x faster |

[ivarflakstad]

Oh yeah big mr huggingface employee coming here showing the people how to vibe code kernel code with proper results, eh?

But seriously this is really good. We've previously talked about how indexing with smaller dtypes is surprisingly efficient, and I see you've also come up with a similar strided index improvement like I did here.
I find this part really interesting (and I suspect is a substantial part of the performance):

if (is_pow2[k]) {
  r = idx & masks[k];
  idx = idx >> shifts[k];
}

and I think it should be unrolled statically like the get_strided_index_1d / get_strided_idx<1> variants.
(edit: at least if all are pow2)

I couldn't find the branch but at some point I had a kernel impl where I calculated the min required index dtype based on the problem size using function constant values. Maybe worth exploring?

[ivarflakstad]

Just checked which kernels are actually being run in the benchmark: fast_sum_f32_strided, fast_sum_f16_strided, and fast_sum_bf16_strided. So we've only hit the IndexType::U32 => "" case and only profiled uint indexing.

[ivarflakstad]

Good news and bad news.

Good news is that I managed to speed this up even more.

Bad news is that I've benchmarked the pow2 concept in isolation and there is no measurable performance benefit.
I'm assuming it is because using 3 lookups (is_pow2[k], masks[k], shifts[k]) counteracts the benefit of mask and shift operations. I even removed the is_pow2[k] check and there was still no measurable difference.
It seems all the gains are from using u32 and static unrolling.

Good news from the bad news - simplifies the improvements considerably!

[ivarflakstad]

Just did some quick mafs.
If the tensor is so large that u32 is not enough to index into it. That means that the underlying buffer has more than 4,294,967,295 elements in it. Even if we were to use fp8 for this tensor it would still be (4294967295 * 8 / (1024^3) =) 32 GB of data in a single tensor. For metal kernels that's kinda silly.
In other words I'm removing the kernels that use u64 to index. We'll use a different approach if the need arises.

[drbh]

Just did some quick mafs. If the tensor is so large that u32 is not enough to index into it. That means that the underlying buffer has more than 4,294,967,295 elements in it. Even if we were to use fp8 for this tensor it would still be (4294967295 * 8 / (1024^3) =) 32 GB of data in a single tensor. For metal kernels that's kinda silly. In other words I'm removing the kernels that use u64 to index. We'll use a different approach if the need arises.

thanks for diving into this! and yea I kinda assumed we could get away with just a u32 but didnt want to make too big of an assumption on todays hardware. however that being said, I do think its reasonable to avoid the u64 for now - and if this become a blocker in the future we can revisit

[drbh]

would approve the refactor/updates from @ivarflakstad if this wasnt originally opened by me!

looks great and agree the standardization on u32 is reasonable and realistic

**note just reran benches and get the following values on a MBP M3

Benchmark	Time	Thrpt
metal_reduce_f32_strided	2.3433 ms	1.6670 GiB/s
metal_reduce_f16_strided	2.3631 ms	846.33 MiB/s
metal_reduce_bf16_strided	2.3545 ms	849.44 MiB/s