#include <cuda/pipeline>

__global__ void example_kernel(const float *in)
{
    constexpr int block_size = 128;
    __shared__ __align__(sizeof(float)) float buffer[4 * block_size];

    // 为每个线程创建一个统一流水线
    cuda::pipeline<cuda::thread_scope_thread> pipeline = cuda::make_pipeline();

    // 内存复制的第一阶段
    pipeline.producer_acquire();
    // 每个线程获取第一个块的一个元素
    cuda::memcpy_async(buffer, in, sizeof(float), pipeline);
    pipeline.producer_commit();

    // 内存复制的第二阶段
    pipeline.producer_acquire();
    // 每个线程获取第二和第三个块的一个元素
    cuda::memcpy_async(buffer + block_size, in + block_size, sizeof(float), pipeline);
    cuda::memcpy_async(buffer + 2 * block_size, in + 2 * block_size, sizeof(float), pipeline);
    pipeline.producer_commit();

    // 内存复制的第三阶段
    pipeline.producer_acquire();
    // 每个线程获取最后一个块的一个元素
    cuda::memcpy_async(buffer + 3 * block_size, in + 3 * block_size, sizeof(float), pipeline);
    pipeline.producer_commit();

    // 等待最旧的阶段（等待第一阶段）
    pipeline.consumer_wait();
    pipeline.consumer_release();

    // __syncthreads();
    // 使用第一阶段的数据

    // 等待最旧的阶段（等待第二阶段）
    pipeline.consumer_wait();
    pipeline.consumer_release();

    // __syncthreads();
    // 使用第二阶段的数据

    // 等待最旧的阶段（等待第三阶段）
    pipeline.consumer_wait();
    pipeline.consumer_release();

    // __syncthreads();
    // 使用第三阶段的数据
}

#include <cuda_pipeline.h>

__global__ void example_kernel(const float *in)
{
    constexpr int block_size = 128;
    __shared__ __align__(sizeof(float)) float buffer[4 * block_size];

    // 第一批内存复制
    // 每个线程获取第一个块的一个元素
    __pipeline_memcpy_async(buffer, in, sizeof(float));
    __pipeline_commit();

    // 第二批内存复制
    // 每个线程获取第二和第三个块的一个元素
    __pipeline_memcpy_async(buffer + block_size, in + block_size, sizeof(float));
    __pipeline_memcpy_async(buffer + 2 * block_size, in + 2 * block_size, sizeof(float));
    __pipeline_commit();

    // 第三批内存复制
    // 每个线程获取最后一个块的一个元素
    __pipeline_memcpy_async(buffer + 3 * block_size, in + 3 * block_size, sizeof(float));
    __pipeline_commit();

    // 等待除最后两批之外的所有内存复制（等待第一批）
    __pipeline_wait_prior(2);

    // __syncthreads();
    // 使用第一批的数据

    // 等待除最后一批之外的所有内存复制（等待第二批）
    __pipeline_wait_prior(1);

    // __syncthreads();
    // 使用第二批的数据

    // 等待所有批次的内存复制（等待第三批）
    __pipeline_wait_prior(0);

    // __syncthreads();
    // 使用最后一批的数据
}

#include <cuda/pipeline>
#include <cooperative_groups.h>

#pragma nv_diag_suppress static_var_with_dynamic_init

using pipeline = cuda::pipeline<cuda::thread_scope_block>;

__device__ void produce(pipeline &pipe, int num_stages, int stage, int num_batches, int batch, float *buffer, int buffer_len, float *in, int N)
{
  if (batch < num_batches)
  {
    pipe.producer_acquire();
    /* copy data from in(batch) to buffer(stage) using asynchronous memory copies */
    pipe.producer_commit();
  }
}

__device__ void consume(pipeline &pipe, int num_stages, int stage, int num_batches, int batch, float *buffer, int buffer_len, float *out, int N)
{
  pipe.consumer_wait();
  /* consume buffer(stage) and update out(batch) */
  pipe.consumer_release();
}

__global__ void producer_consumer_pattern(float *in, float *out, int N, int buffer_len)
{
  auto block = cooperative_groups::this_thread_block();

  /* Shared memory buffer declared below is of size 2 * buffer_len
     so that we can alternatively work between two buffers.
     buffer_0 = buffer and buffer_1 = buffer + buffer_len */
  __shared__ extern float buffer[];

  const int num_batches = N / buffer_len;

  // Create a partitioned pipeline with 2 stages where half the threads are producers and the other half are consumers.
  constexpr auto scope = cuda::thread_scope_block;
  constexpr int num_stages = 2;
  cuda::std::size_t producer_count = block.size() / 2;
  __shared__ cuda::pipeline_shared_state<scope, num_stages> shared_state;
  pipeline pipe = cuda::make_pipeline(block, &shared_state, producer_count);

  // Fill the pipeline
  if (block.thread_rank() < producer_count)
  {
    for (int s = 0; s < num_stages; ++s)
    {
      produce(pipe, num_stages, s, num_batches, s, buffer, buffer_len, in, N);
    }
  }

  // Process the batches
  int stage = 0;
  for (size_t b = 0; b < num_batches; ++b)
  {
    if (block.thread_rank() < producer_count)
    {
      // Prefetch the next batch
      produce(pipe, num_stages, stage, num_batches, b + num_stages, buffer, buffer_len, in, N);
    }
    else
    {
      // Consume the oldest batch
      consume(pipe, num_stages, stage, num_batches, b, buffer, buffer_len, out, N);
    }
    stage = (stage + 1) % num_stages;
  }
}