[CINN] Add the ReplaceCrossBlockReduction backend pass

[lshpku]

PR Category

CINN

PR Types

New features

Description

This PR adds the ReplaceCrossBlockReduction pass in optim::Optimize, and provides the CUDA implementations for grid_reduce and update_semaphore.

The ReplaceCrossBlockReduction pass performs the following tasks:

• Replaces the cross-block reduction with an external call to grid_reduce.
• Adds a condition check is_last_block_done to the reduction operation and all subsequent schedule blocks.
• Adds global buffers (rf and semaphore) to the function’s argument list.

Example:

function reduce_sum (..., var_1)
{
  thread_bind[blockIdx.x] for (i, 0, 16):
    thread_bind[blockIdx.y] for (j, 0, 8): // reduce axis
      var_1[i] += var_1_rf[j, i]
}

After pass:

function reduce_sum (..., var_1, var_1_rf, semaphore)
{
  thread_bind[blockIdx.x] for (i, 0, 16):
    thread_bind[blockIdx.y] for (j, 0, 8): // reduce axis
      is_last_block_done = update_semaphore(semaphore)
      if (is_last_block_done):
        var_1[i] = grid_reduce_sum(var_1_rf)
}

Pcard-85711

[paddle-ci-bot]

Sorry to inform you that 6b78d4b's CIs have passed for more than 7 days. To prevent PR conflicts, you need to re-run all CIs manually.