Use new GPU kernel for [unsorted] segment reductions

[benbarsdell]

• Optionally replaces the old atomics-based kernels with calls to SegmentReduceGPU (the same kernel already used for sparse segment reductions). This behavior is enabled by default, but the old kernels can be re-enabled by setting the environment variable TF_USE_ATOMIC_SEGMENT_REDUCTIONS=1. On Windows, the old kernels are always used due to a build issue with the new kernel.
• This improves performance, and guarantees that these ops are deterministic. In future it is hoped that the old kernels can be removed completely.
• Also adds a GPU kernel registration for SegmentMean, which didn't previously exist.

cc @nluehr @reedwm @duncanriach

[sanjoy]

Can we call this UseNonDeterministicSegmentReductions?

[benbarsdell]

Done

[sanjoy]

Can you please add /*EmptySegmentValueF=*/ and /*InitialValueF=*/ here to make it clear which arg is which? Same elsewhere.

[benbarsdell]

Done

[duncanriach]

First pass through this. Nice work, @benbarsdell. A couple of thoughts:

1. The original code differentiated between atomic and nonatomic reductions (because there were sometimes both), but not for reasons of determinism. And, of course, atomics can be used in deterministic code. So, there's some lingering old symbology around atomic/nonatomic which has gotten tangled up with (or repurposed for) deterministic/nondeterministic operation. @sanjoy picked out one example of that above. There are other places in the code where perhaps "non-atomic" would ideally be replaced with "deterministic" and where "atomic" would ideally be replaced with "nondeterministic."
2. I'm wondering if we should add tests to confirm determinism, or if we should rely on it being deterministic by design. There will, of course, be backup testing from @reedwm's determinism auto-checker.

[benbarsdell]

Thanks for the suggestions Duncan. I'm facing some build issues right now but will push the additional atomic->deterministic rewordings tomorrow along with a rebase.

I haven't looked closely at how the determinism tests work but will take a look.

[reedwm]

This was rolled back in 44cdcda. Will investigate.

[reedwm]

Rolled forward in 9a1072e

[reedwm]

Unfortunately this was rolled back again in 51ee415, since it caused a performance regression in an internal model. Also, the test segment_reduction_ops_deterministic_test.py failed on Windows, since the deterministic algorithms were disabled on Windows.

I think the best approach here is to keep using the old nondeterministic kernels by default, and only use the new ones if either determinism is enabled or an environmental variable is set (say, TF_USE_DETERMINISTIC_SEGMENT_REDUCTIONS). Then the determinism API is unblocked and we can internally debug the performance issue, and potentially send you an example to reproduce.

@benbarsdell, do you want to create a new PR with the new kernels disabled by default, by this Wednesday? If not, I can rollforward, disabling the new kernels by default. I want to get this in as soon as possible since it is required for determinism.