[webgpu] Fix atomic shared mem load inside loop

[wpmed92]

test_arange::TestIndexing::test_index_mnist_opt with GROUPTOP causes a shared mem load inside a loop, which didn't get rewritten to packed load because of a missing allow_any_len=True in the wgsl pattern matcher.

The load inside loop part of the problematic kernel before the change:

if (((bool(lidx0))!=true)) {
    var acc0 = 0u;
    for (var ridx1001 = 0; ridx1001 < 16; ridx1001++) {
      var val4 = atomicLoad(&temp0[ridx1001]);
      acc0 = (acc0+val4);
    }
}

after the change with proper packed load.

if (((bool(lidx0))!=true)) {
    var acc0 = 0u;
    for (var ridx1001 = 0; ridx1001 < 16; ridx1001++) {
      var val4 = atomicLoad(&temp0[(ridx1001>>2)]);
      acc0 = (acc0+(u32(((val4>>(((u32(ridx1001))&3u)<<3u))&255u))));
    }
}

Although the load was atomic, we can see that the uchar is not shifted out in the first version.

[chenyuxyz]

do you know why atomic is wrong for webgpu shared memory?

[wpmed92]

I haven't yet find out, i know that the packed load/store is the same as for "normal" memory from a webgpu codegen pov, so there's no reason it should give bad result.
I'll check with wgpu-py to see if it's a dawn issue and in the meantime I write some wgsl kernels that trigger this.

[github-actions]

Changes

Name                         Lines    Diff    Tokens/Line    Diff
-------------------------  -------  ------  -------------  ------
tinygrad/renderer/wgsl.py       82      -1           23.0    +0.2


total lines changes: -1

[wpmed92]

@chenyuxyz I did my tests and it turned out I was wrong regarding shared mem atomics being wrong in webgpu. It was a missing allow_any_len=True that caused the missed packed load rewrite. I updated the readme of the PR.

[chenyuxyz]

cool yea this looks more like the true fix