Support GPU-based reductions

[e-kayrakli]

I am filing this on behalf of a user. We lacked reduction support on GPUs so far, mainly because it got de-prioritized over other user-driven requests. But I think we can start working on reductions soon. Here's how I think we should tackle this:

Step 1: Provide reduction functions for simple cases to move forward.

We can think of gpuSumReduce(x: []) et al. that can be used as a drop-in replacement for simple + reduce x where x is on the GPU. We can make gpuSumReduce(x: []) portable between the host and the device. I imagine the device side of the implementation can use CUB for NVIDIA and its analogue (hoping that there's one) for AMD. Essentially gpuSumReduce will be a wrapper for CUB's reduction support. That's probably as best as we can get in terms of performance. Internal note for this: https://github.com/Cray/chapel-private/issues/5330

Optional Step 2: Lower + reduce x into gpuSumReduce

We can imagine lowering such reductions into our GPU support function. Otherwise, I believe that will lower into a forall with reduce intent. I think as long as the non-gpu branch in gpuSumReduce is still + reduce x, we can probably pull this off with no performance impact.

We can also do this with GPU specialization (--gpu-specialization)

Step 3: Get reduce intents to work

This is our ultimate goal. It depends on #18500, which is a big lift in its own right, but it is already in the works.

[milthorpe]

I understand from this issue that reductions on GPUs are not supported in Chapel 1.32. As this isn't well documented in the GPU programming technote, I wanted to document an interesting / weird behaviour I noticed with forall with a reduce intent on GPU. My test program is as follows:

proc computeSum(loc) {
    on loc {
        var sum: int = 0;
        forall a in 0..#10 with (+ reduce sum) {
            sum += a;
            writeln(sum);
        }
        writeln(sum);
    }
}

if here.gpus.size > 0 then computeSum(here.gpus[0]);
else computeSum(here);

On CPU, as expected, the final number printed is 45. (As forall intents accumulate to a per-task reduction variable, I believe the output of the writeln statement inside the forall loop is undefined. On my system, it prints the numbers 0..9 in random order.)
On GPU, the final number printed by the code above is also 45, however, if I remove the first writeln statement from inside the loop, the final number printed is now 0.

It was surprising to me that simply printing the intermediate results would change the final result of the computation. I assume this behaviour will disappear in Chapel 1.33, but I wanted to record it here in case other users hit the same issue.

[bradcray]

@e-kayrakli can probably say much more authoritatively, but my guess is that the writeln() causes the loop not to be executed on the GPU which is why it works. Meanwhile, when it's executed on the GPU, I'm guessing we're not implementing the reduce intent properly (which wouldn't be completely surprising since GPUs are driven by foreach loops and foreach loops don't have good with-clause support yet... though we're working on that.

In any case, it'd be nice if cases like this would generate a compiler error rather than the wrong answer, at least until they're working properly. I suspect you've just poked into a corner that others haven't yet—thanks for pointing it out, Josh. Would you be up for creating a bug issue for this?

[milthorpe]

I also note that if I add @assertOnGPU to the forall loop, I do not get a compile error with CHPL_LOCALE_MODEL=gpu, i.e. it looks to me as a user that the loop compiled correctly for GPU.
I'm happy to create a bug issue along the lines of 'forall reductions fail silently for GPU'; as you suggest, it would be nicer if they noisily failed to compile.

[bradcray]

I also note that if I add @assertOnGPU to the forall loop, I do not get a compile error with CHPL_LOCALE_MODEL=gpu

That's with the writeln() in the loop? That definitely seems surprising/confusing...

[e-kayrakli]

Brad's assessment sounded right to me. I thought we had a issue capturing this wrong result behavior but it is from way past and it is in our private tracker.

That's with the writeln() in the loop? That definitely seems surprising/confusing...

I agree, I'll take a quick look.

[e-kayrakli]

I also note that if I add @assertOnGPU to the forall loop, I do not get a compile error with CHPL_LOCALE_MODEL=gpu, i.e. it looks to me as a user that the loop compiled correctly for GPU.

I can't reproduce this on main.

proc computeSum(loc) {
  on loc {
    var sum: int = 0;
    @assertOnGpu
    forall a in 0..#10 with (+ reduce sum) {
      sum += a;
      writeln(sum);
    }
    writeln(sum);
  }
}

if here.gpus.size > 0 then computeSum(here.gpus[0]);
else computeSum(here);

Produces:

warning: The prototype GPU support implies --no-checks. This may impact debuggability. To suppress this warning, compile with --no-checks explicitly
$CHPL_HOME/forallReduce.chpl:1: In function 'computeSum':
$CHPL_HOME/forallReduce.chpl:5: error: Loop is marked with @assertOnGpu but is not eligible for execution on a GPU
$CHPL_HOME/forallReduce.chpl:7: note: call has outer var access
  $CHPL_HOME/forallReduce.chpl:13: called as computeSum(loc: locale)
note: generic instantiations are underlined in the above callstack

[milthorpe]

Sorry - you're right - I tried it on a version of the code where I'd already removed the writeln. With the writeln, I get the same error Engin observed above.

[e-kayrakli]

This is done via #24787

Follow ups are in #24932