Fix normExp shared-memory reductions and odd-rollout handling

[plancherb1]

We were trying to use this project and noticed some small bugs that we
wanted to upstream in case others run into similar issues. This change
was AI-assisted, reviewed by the authors, and validated using the tests
reported below.

-Brian

computeBaselineCost and computeNormalizer in
include/mppi/core/mppi_common.cu have two defects that make the
single-block weight-update path
(fullGPUcomputeWeights / launchWeightTransformKernel) return
incorrect baselines and normalizers:

1. __restrict__ on the multi-writer shared reduction buffer. Both
   helpers take reduction_buffer as float* __restrict__, but that
   buffer is written by peer threads across __syncthreads(). The
   qualifier is inappropriate for a buffer with cross-thread visibility
   requirements. On an RTX 5090 with CUDA 13.2, the restricted version
   produces incorrect results and removing the qualifier resolves them:
   the existing
   NormExpKernel.comparisonTestHostvsDeviceBaselineNormalizerCalculation
   test fails on current main (device baseline is not the true minimum
   at 10,000 rollouts, data-dependently). A small qualifier matrix
   isolates the shared-buffer qualification as the cause —
   float* __restrict__ fails; plain float*,
   volatile float* __restrict__, and volatile float* all pass. We
   removed the qualifier rather than adding volatile:
   __syncthreads() provides the required block-wide memory ordering,
   and CUDA's documentation is explicit that volatile is not a
   synchronization mechanism.

2. Odd-rollout leftover handling reads unwritten shared memory, and
   races. With odd num_rollouts, the stage-1 fixup folds in
   reduction_buffer[num_rollouts - 1], but stage 1 only writes
   reduction_buffer[0 .. num_rollouts/2), so that slot is
   uninitialized (the leftover element actually lives in
   trajectory_costs_d[num_rollouts - 1]); the fixup also reads
   reduction_buffer[prev_size - 1] from a different thread with no
   intervening barrier. Fixed by folding the leftover element into the
   first-stage loop, in the same thread that writes
   reduction_buffer[prev_size - 1] — no extra barrier needed. The
   later in-loop fixup is correct and unchanged.

Tests:

• New odd-rollout device cases for both helpers at N=999 and N=6049
  (below and above the 1024-thread launch). The baseline case places the
  true minimum at the last index — the exact element the leftover
  handling is responsible for; the normalizer case uses exactly
  representable integer values so the serial reference and parallel tree
  sum are bit-identical and the exact-equality assert is independent of
  reduction order. Both fail on current main (e.g. baseline 9
  vs. true minimum 1; sum off by 4x).
• A focused fullGPUcomputeWeightsEveryIteration_Test at 10,000
  rollouts that samples fresh costs and checks every iteration with a
  SCOPED_TRACE naming the failing dataset (the existing comparison
  test only asserts the final iteration, and the defect is
  data-dependent). On current main it fails at iteration 1; with the
  fix it passes all 500 iterations.
• The existing test-local wrapper kernels had every thread writing the
  scalar *output; the store is now gated on threadIdx.x == 0.

Validation on an RTX 5090 (CUDA 13.2): full normexp_kernel_tests
passes 10/10 with the fix, including the previously failing
host-vs-device comparison test; rmppi_kernel_tests passes 5/5;
compute-sanitizer --tool racecheck over all device-touching tests in
the suite reports 0 hazards.

[bogidude]

Hi Brian,

Thank you for the PR! There are definitely some bugs in that section of the code as I haven't looked at it in years. At some point, I tried to do the normalizer and baseline computations on the GPU and found them to be slower than running them on the CPU (at the time anyway) so I never added it to any controller's critical path. And because that code is not used by anything except unit tests, it has been left untouched and buggy. I'm actually surprised that the unit tests have passed on the multiple different computers I have tried them on in the past upon reading your PR.

I do have some remarks about the proposed fixes:

1. You are saying that the __restrict__ keyword is the cause of the issues with miscomputed values but I think I disagree. The __restrict__ keyword should be a hint to the compiler that __restrict__ pointers don't overlap in memory (i.e. accessing any value in array A won't overlap with any items in array B). The actual bug looks to be a race condition where thread A would write to reduction_buffer[i] and then also possibly try to access items reduction_buffer[num_rollouts - 1] and reduction_buffer[prev_size - 1] if thread A was the one entering the if statement. There was no __syncthreads() between writing to reduction_buffer[i] and reading from reduction_buffer[num_rollouts - 1] and reduction_buffer[prev_size - 1] so a race condition was produced if the threads were resolved in different orders. Your code looks to fix that by utilizing a temp variable val and reads only from trajectory_costs_d and only writes to reduction_buffer. This prevents the need for __syncthreads() at the cost of reading from global memory one more time within the if statement. Whether utilizing a __syncthreads() vs adding an additional read to global memory is faster is going to be depend on the actual use case but I think the vast majority of the time, your changes would be faster (__syncthreads() is so slow in my experience). But the discussion of the __restrict__ keyword affecting this seems completely wrong to me. Could you try adding back the __restrict__ keyword with the rest of your changes to see if it still causes problems for you? I haven't messed with CUDA 13 or an RTX 5090 so it could be that __restrict__ causes additional issues but I don't see why it would in this case.
2. Yeah, when I wrote this code, I definitely wasn't thinking about odd numbers of rollouts so this is a genuine bug that would need fixing. The code changes you are suggesting should address the issue from my understanding. There is some similar code that is actually used multiple times in all the MPPI-Generic controllers in costArrayReduction(). I am pretty sure it can handle odd-sized arrays and race conditions properly but I wouldn't mind a second set of eyes looking over it to double check if you have the time.

As this code is currently not used anywhere in the library but the unit tests, I am curious how you stumbled upon the issue in the first place. Was it just the unit test failing or are you trying to use the fullGPUcomputeWeights / launchWeightTransformKernel methods in your own custom controller?

If you can get back to me about the result of adding the __restrict__ keyword back, I will be happy to merge this in! If it is still causing issues for you, that's a concerning result that might affect other parts of the library and potentially require a lot more rewriting on my part.