[BUG]: block-level InclusiveScan with initial value doesn't compute the block aggregate correctly

[Beanavil]

Is this a duplicate?

• I confirmed there appear to be no duplicate issues for this bug and that I agree to the Code of Conduct

Type of Bug

Silent Failure

Component

CUB

Describe the bug

The recently-added block-level InclusiveScan with initial value support seems to not be taking into account such initial value for the block aggregate calculation.

How to Reproduce

1. Add the following test

__global__ void InclusiveScanKernelAggregate(int* output, int* d_block_aggregate)
{
  // Specialize BlockScan for a 1D block of 64 threads of type int
  using block_scan_t   = cub::BlockScan<int, 64>;
  using temp_storage_t = block_scan_t::TempStorage;

  // Allocate shared memory for BlockScan
  __shared__ temp_storage_t temp_storage;

  int initial_value = 1;
  int thread_data[] = { 1, -1 };
  //  input: {[1, -1], [1, -1],[1, -1], ... [1, -1]}

  // Collectively compute the block-wide inclusive scan max
  int block_aggregate;
  block_scan_t(temp_storage).InclusiveScan(thread_data, thread_data, initial_value, cub::Sum(), block_aggregate);

  // output: {[2, 1], [2, 1],[2, 1], ... [2, 1]}
  // block_aggregate = 1;
  // ...
  // example-end inclusive-scan-array-aggregate-init-value

  *d_block_aggregate          = block_aggregate;
  output[threadIdx.x * 2]     = thread_data[0];
  output[threadIdx.x * 2 + 1] = thread_data[1];
}

CUB_TEST("Block array-based inclusive scan with block aggregate works with initial value", "[scan][block]")
{
  thrust::device_vector<int> d_out(block_num_threads * num_items_per_thread);

  c2h::device_vector<int> d_block_aggregate(1);
  InclusiveScanKernelAggregate<<<1, block_num_threads>>>(
    thrust::raw_pointer_cast(d_out.data()), thrust::raw_pointer_cast(d_block_aggregate.data()));
  REQUIRE(cudaSuccess == cudaPeekAtLastError());
  REQUIRE(cudaSuccess == cudaDeviceSynchronize());

  c2h::host_vector<int> expected(d_out.size());
  for (size_t i = 0; i < expected.size() - 1; i += 2)
  {
    expected[i]     = 2;
    expected[i + 1] = 1;
  }

  REQUIRE(d_out == expected);
  REQUIRE(d_block_aggregate[0] == 1);
}

to catch2_test_block_scan_api.cu.

2. Build and run the test. E.g:

./ci/build_cub.sh

cmake --build build/cub-cpp17/ --target cub.cpp17.test.block_scan_api

./build/cub-cpp17/cub.cpp17.test.block_scan_api

It should result in a failure due to the block aggregate computed being 0 and not 1.

FYI: a similar test is implemented in catch2_test_block_scan_api.cu, but it uses the cub::Max() operator so it does seem to oversee the aggregates miscalculation.

Expected behavior

The block InclusiveScan with initial values should also include the initial value in the aggregates computation.

Reproduction link

No response

Operating System

No response

nvidia-smi output

No response

NVCC version

No response