pufferlib.cu

#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

namespace pufferlib {

__host__ __device__ void puff_advantage_row_cuda_fallback(float* values, float* rewards, float* dones,
        float* importance, float* advantages, float gamma, float lambda,
        float rho_clip, float c_clip, int horizon) {
    float lastpufferlam = 0;
    for (int t = horizon-2; t >= 0; t--) {
        int t_next = t + 1;
        float nextnonterminal = 1.0 - dones[t_next];
        float rho_t = fminf(importance[t], rho_clip);
        float c_t = fminf(importance[t], c_clip);
        float delta = rho_t*(rewards[t_next] + gamma*values[t_next]*nextnonterminal - values[t]);
        lastpufferlam = delta + gamma*lambda*c_t*lastpufferlam*nextnonterminal;
        advantages[t] = lastpufferlam;
    }
}


__device__ __forceinline__ void puff_advantage_row_cuda(float* values, float* rewards, float* dones,
        float* importance, float* advantages, float gamma, float lambda,
        float rho_clip, float c_clip, int horizon) {
    
    // Fall back to original if horizon not divisible by 4
    if (horizon % 4 != 0) {
        puff_advantage_row_cuda_fallback(values, rewards, dones,
            importance, advantages, gamma, lambda, rho_clip, c_clip, horizon);
        return;
    }
    
    float lastpufferlam = 0.0f;
    int num_chunks = horizon / 4;
    
    // need to track values across chunks
    float next_value = values[horizon - 1];
    float next_done = dones[horizon - 1];
    float next_reward = rewards[horizon - 1];
    
    // Process chunks from end to beginning
    for (int chunk = num_chunks - 1; chunk >= 0; chunk--) {
        int base = chunk * 4;
        
        // Load 4 elements at once
        float4 v4 = *reinterpret_cast<float4*>(values + base);
        float4 r4 = *reinterpret_cast<float4*>(rewards + base);
        float4 d4 = *reinterpret_cast<float4*>(dones + base);
        float4 i4 = *reinterpret_cast<float4*>(importance + base);
        
        float v[4] = {v4.x, v4.y, v4.z, v4.w};
        float r[4] = {r4.x, r4.y, r4.z, r4.w};
        float d[4] = {d4.x, d4.y, d4.z, d4.w};
        float imp[4] = {i4.x, i4.y, i4.z, i4.w};
        float adv[4] = {0.0f, 0.0f, 0.0f, 0.0f};
        
        int start_idx = (chunk == num_chunks - 1) ? 2 : 3;
        
        #pragma unroll
        for (int i = start_idx; i >= 0; i--) {
            float nextnonterminal = 1.0f - next_done;
            float rho_t = fminf(imp[i], rho_clip);
            float c_t = fminf(imp[i], c_clip);
            float delta = rho_t * (next_reward + gamma * next_value * nextnonterminal - v[i]);
            lastpufferlam = delta + gamma * lambda * c_t * lastpufferlam * nextnonterminal;
            adv[i] = lastpufferlam;
            
            next_value = v[i];
            next_done = d[i];
            next_reward = r[i];
        }
        
        float4 adv4 = make_float4(adv[0], adv[1], adv[2], adv[3]);
        *reinterpret_cast<float4*>(advantages + base) = adv4;
    }
}

void vtrace_check_cuda(torch::Tensor values, torch::Tensor rewards,
        torch::Tensor dones, torch::Tensor importance, torch::Tensor advantages,
        int num_steps, int horizon) {

    // Validate input tensors
    torch::Device device = values.device();
    for (const torch::Tensor& t : {values, rewards, dones, importance, advantages}) {
        TORCH_CHECK(t.dim() == 2, "Tensor must be 2D");
        TORCH_CHECK(t.device() == device, "All tensors must be on same device");
        TORCH_CHECK(t.size(0) == num_steps, "First dimension must match num_steps");
        TORCH_CHECK(t.size(1) == horizon, "Second dimension must match horizon");
        TORCH_CHECK(t.dtype() == torch::kFloat32, "All tensors must be float32");
        if (!t.is_contiguous()) {
            t.contiguous();
        }
    }
}

 // [num_steps, horizon]
__global__ void puff_advantage_kernel(float* values, float* rewards,
        float* dones, float* importance, float* advantages, float gamma,
        float lambda, float rho_clip, float c_clip, int num_steps, int horizon) {
    int row = blockIdx.x*blockDim.x + threadIdx.x;
    if (row >= num_steps) {
        return;
    }
    int offset = row*horizon;
    puff_advantage_row_cuda(values + offset, rewards + offset, dones + offset,
        importance + offset, advantages + offset, gamma, lambda, rho_clip, c_clip, horizon);
}

void compute_puff_advantage_cuda(torch::Tensor values, torch::Tensor rewards,
        torch::Tensor dones, torch::Tensor importance, torch::Tensor advantages,
        double gamma, double lambda, double rho_clip, double c_clip) {
    int num_steps = values.size(0);
    int horizon = values.size(1);
    vtrace_check_cuda(values, rewards, dones, importance, advantages, num_steps, horizon);
    TORCH_CHECK(values.is_cuda(), "All tensors must be on GPU");

    int threads_per_block = 32;
    int blocks = (num_steps + threads_per_block - 1) / threads_per_block;

    puff_advantage_kernel<<<blocks, threads_per_block>>>(
        values.data_ptr<float>(),
        rewards.data_ptr<float>(),
        dones.data_ptr<float>(),
        importance.data_ptr<float>(),
        advantages.data_ptr<float>(),
        gamma,
        lambda,
        rho_clip,
        c_clip,
        num_steps,
        horizon
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        throw std::runtime_error(cudaGetErrorString(err));
    }
}

TORCH_LIBRARY_IMPL(pufferlib, CUDA, m) {
  m.impl("compute_puff_advantage", &compute_puff_advantage_cuda);
}

}