[SYCL] Support shuffle algorithms for non-uniform groups

[steffenlarsen]

This commit makes the non-uniform group classes support the shift and select algorithms.

[Pennycook]

Thank you for working on this, and for tidying up some of the SPIR-V wrappers!

[Pennycook]

The changes look good to me, modulo the test failures.

[steffenlarsen]

@JackAKirk - Can you think of a reason why this would fail for the new test cases for ballot_group only?

[JackAKirk]

I think it is failing because the delta does not relate to the true delta only considering set bits in the mask which is what is desired. I think could also fail for opportunistic_group, but guess this depends on the test/execution.

I think for these cases (opportunistic_group and ballot_group, but you could also use this for all non-uniform groups) if you replace

    return __nvvm_shfl_sync_up_i32(detail::ExtractMask(detail::GetMask(g))[0],
                                   x, delta, 0);

with a call to

non_uniform_shfl

defined here:

https://github.com/intel/llvm/blob/sycl/sycl/include/sycl/ext/oneapi/experimental/cuda/non_uniform_algorithms.hpp#L151

as e.g.

unsigned localSetBit = g.get_local_id()[0] + 1;
int unfoldedSrcSetBit = localSetBit + delta;
auto MemberMask = detail::ExtractMask(detail::GetMask(g))[0];
return non_uniform_shfl(g, MemberMask, x,
                                __nvvm_fns(MemberMask, 0, unfoldedSrcSetBit));

Then it should work.

[JackAKirk]

Actually thinking about it more, in this case I think the sign of delta might have to be changed for ballot_group/opportunistic group, when using that non_uniform_shfl function.
e.g. in my above message you might have to replace

int unfoldedSrcSetBit = localSetBit + delta;

with:

int unfoldedSrcSetBit = localSetBit - delta;

In any case I think one of those should work. I can't remember the docs offhand for which version is correct (basically whether the semantic is send to idx or receive from). I think it is the - delta version!

[steffenlarsen]

Thanks a ton, Jack! Let's give it a try! 😄

[steffenlarsen]

Sadly it doesn't seem that it changed anything, but it also doesn't seem like we call non_uniform_shfl here. I wonder if it is immediately usable or if it has been specialized for the use-case with reduce and scan. Either way, I do not have an easy way of testing fixes, so would you be okay with me disabling the subset of checks for CUDA and opening an issue so you can have a look when possible? (Tag @npmiller )

[JackAKirk]

Sure, no worries. I think I know how to fix it.

[steffenlarsen]

@intel/llvm-reviewers-runtime @sergey-semenov @uditagarwal97 - Friendly ping.

[uditagarwal97]

Within the bounds of my limited knowledge in this part of the code base, the changes LGTM!

[steffenlarsen]

@intel/llvm-reviewers-runtime | @sergey-semenov - Friendly ping.