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 is due to an implicit assumption in subgroup and workgroup mapping. For example the current topk_moe shader assumes that each subgroup lane is contiguous in the x dimension of a workgroup which is not exactly guaranteed .

• There is no direct relationship between SubgroupLocalInvocationId and LocalInvocationId or LocalInvocationIndex

On the failing platform we were observing a single subgroup getting mapped in a 2D workgroup ID. For subgroupID 0 we get gl_LocalInvocationID.x: 0-7 and gl_LocalInvocationID.y: 0-3, whereas on other platforms it was gl_LocalInvocationID.x: 0-31 and gl_LocalInvocationID.y: 0.

We've tested 4404be4 on Intel/NVIDIA/AMD GPU for Windows, Intel/NVIDIA GPU for Linux and confirmed test-backend-ops.exe -o TOPK_MOE passes

[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.