vulkan: Fix mismatch in TOPK_MOE unit test

[rillomas]

We are currently seeing a mismatch in all TOPK_MOE unit tests on our upcoming platform. This seems to be due to an implicit assumption in subgroup and workgroup mapping. For example the current shader assumes that each subgroup lane is contiguous in the x dimension of a workgroup which is not exactly guaranteed (though this may not be the exact reason for the mismatch we're seeing now).

Currently experimenting for the fix.

[jeffbolznv]

I agree a change like this is necessary due to this spec language: "There is no direct relationship between SubgroupLocalInvocationId and LocalInvocationId or LocalInvocationIndex". I guess it's a problem specifically for this shader because it (1) uses local_size_y > 1 and (2) uses subgroup instructions?

I wonder if a similar problem could be affecting the shader in #17389, though it doesn't use local_size_y > 1 so maybe not.

[rillomas]

Thanks for your comment. Yes, this may be one of those issues that were hidden due to the driver's default behavior.

[rillomas]

It turns out that a single subgroup was mapped in a 2D workgroup on our failing environment. For subgroupID 0 we get gl_LocalInvocationID.x: 0-7 and gl_LocalInvocationID.y: 0-3, whereas previously it was gl_LocalInvocationID.x: 0-31 and gl_LocalInvocationID.y: 0. 62ed513 passes all test-backend-ops on the target environment but still need to cleanup and test other environments.