turn 0675.1s turn · 85.1s LLM · 590.0s tools · tokens in=1,430 · out=2,835 · total=4,728 · FINAL
reasoning
The task is to replace the forward computation of a PyTorch model with a custom CUDA kernel. The model is a Hamiltonian Monte Carlo (HMC) proposal using leapfrog integration.
I need to create a ModelNew class that replaces the forward computation of a PyTorch model with a CUDA kernel for faster HMC proposal using leapfrog integration. The output is a complete Python file using torch.utils.cpp_extension.load_inline to compile the CUDA code.
assistant```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <curand_kernel.h>
__global__ void hmc_leapfrog_kernel(
const float* __restrict__ q_in,
float* __restrict__ q_out,
bool* __restrict__ accepted_out,
const int batch_size,
const int dim,
const float eps,
const int num_leapfrog_steps
) {
const int b = blockIdx.x * blockDim.x + threadIdx.x;
if (b >= batch_size) return;
extern __shared__ float shared_mem[];
float* p = shared_mem + threadIdx.x * dim;
float* q = shared_mem + (blockDim.x * dim) + threadIdx.x * dim;
float* current_q = shared_mem + (2 * blockDim.x * dim) + threadIdx.x * dim;
float* current_p = shared_mem + (3 * blockDim.x * dim) + threadIdx.x * dim;
// Load q
for (int d = 0; d < dim; d++) {
q[d] = q_in[b * dim + d];
current_q[d] = q[d];
}
// Sample p ~ N(0,1) - use simple LCG for speed, or curand
curandState state;
curand_init(1234ULL + b, 0, 0, &state);
float sum_p2 = 0.0f;
for (int d = 0; d < dim; d++) {
float pd = curand_normal(&state);
p[d] = pd;
current_p[d] = pd;
sum_p2 += pd * pd;
}
// Half step momentum: p += -0.5 * eps * q
float half_eps = 0.5f * eps;
for (int d = 0; d < dim; d++) {
p[d] -= half_eps * q[d];
}
// Leapfrog steps
for (int step = 0; step < num_leapfrog_steps - 1; step++) {
// q += eps * p
for (int d = 0; d < dim; d++) {
q[d] += eps * p[d];
}
// p += eps * (-q)
for (int d = 0; d < dim; d++) {
p[d] -= eps * q[d];
}
}
// Final full step
for (int d = 0; d < dim; d++) {
q[d] += eps * p[d];
}
for (int d = 0; d < dim; d++) {
p[d] -= half_eps * q[d];
}
// p = -p
float sum_p2_new = 0.0f;
for (int d = 0; d < dim; d++) {
p[d] = -p[d];
sum_p2_new += p[d] * p[d];
}
// Compute energies
float current_H = 0.0f;
float proposed_H = 0.0f;
for (int d = 0; d < dim; d++) {
current_H += 0.5f * current_q[d] * current_q[d];
proposed_H += 0.5f * q[d] * q[d];
}
current_H += 0.5f * sum_p2;
proposed_H += 0.5f * sum_p2_new;
float log_alpha = current_H - proposed_H;
// Random
float log_u = logf(curand_uniform(&state) + 1e-10f);
bool acc = log_u < log_alpha;
// Write output
for (int d = 0; d < dim; d++) {
q_out[b * dim + d] = acc ? q[d] : current_q[d];
}
accepted_out[b] = acc;
}
torch::Tensor hmc_forward_cuda(
torch::Tensor q,
float step_size,
int num_leapfrog_steps,
int dim
) {
const int batch_size = q.size(0);
auto q_out = torch::empty_like(q);
auto accepted = torch::empty({batch_size}, q.options().dtype(torch::kBool));
const int threads = 128;
const int blocks = (batch_size + threads - 1) / threads;
size_t shared_mem_size = threads * 4 * dim * sizeof(float);
hmc_leapfrog_kernel<<<blocks, threads, shared_mem_size>>>(
q.data_ptr<float>(),
q_out.data_ptr<float>(),
accepted.data_ptr<bool>(),
batch_size,
dim,
step_size,
num_leapfrog_steps
);
return q_out;
}
std::vector<torch::Tensor> hmc_forward_full(
torch::Tensor q,
float step_size,
int num_leapfrog_steps,
int dim
) {
const int batch_size = q.size(0);
auto q_out = torch::empty_like(q);
auto accepted = torch::empty({batch_size}, q.options().dtype(torch::kBool));
const int threads = 128;
const int blocks = (batch_size + threads - 1) / threads;
size_t shared_mem_size = threads * 4 * dim * sizeof(float);
hmc_leapfrog_kernel<<<blocks, threads, shared_mem_size>>>(
q.data_ptr<float>(),
q_out.data_ptr<float>(),
accepted.data_ptr<bool>(),
batch_size,
dim,
step_size,
num_leapfrog_steps
);
return {q_out, accepted};
}
'''
cpp_source = r'''
#include <torch/extension.h>
std::vector<torch::Tensor> hmc_forward_full(torch::Tensor q, float step_size, int num_leapfrog_steps, int dim);
torch::Tensor hmc_forward_cuda(torch::Tensor q, float step_size, int num_leapfrog_steps, int dim);
'''
module = load_inline(
name='hmc_cuda',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['hmc_forward_full'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self, dim, num_leapfrog_steps=10, step_size=0.1):
super().__init__()
self.dim = dim
self.num_leapfrog_steps = num_leapfrog_steps
self.step_size = step_size
self.register_buffer("mu", torch.zeros(dim))
def forward(self, q: torch.Tensor) -> tuple:
results = module.hmc_forward_full(q.contiguous(), self.step_size, self.num_leapfrog_steps, self.dim)
return results[0], results[1]
```
The kernel is submitted. It uses shared memory per block for the state and curand for RNG. This should be faster than the PyTorch reference for the given sizes. The return is q_new and accepted as required. Numerical equivalence should hold within fp32 tolerance.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <curand_kernel.h>\n\n__global__ void hmc_leapfrog_kernel(\n const float* __restrict__ q_in,\n float* __restrict__ q_out,\n bool* __restrict__ accepted_out,\n const int batch_size,\n const int dim,\n const float eps,\n const int num_leapfrog_steps\n) {\n const int b = blockIdx.x * blockDim.x + threadIdx.x;\n if (b >= batch_size) return;\n\n extern __shared__ float shared_mem[];\n float* p = shared_mem + threadIdx.x * dim;\n float* q = shared_mem + (blockDim.x * dim) + threadIdx.x * dim;\n float* current_q = shared_mem + (2 * blockDim.x * dim) + threadIdx.x * dim;\n float* current_p = shared_mem + (3 * blockDim.x * dim) + threadIdx.x * dim;\n\n // Load q\n for (int d = 0; d < dim; d++) {\n q[d] = q_in[b * dim + d];\n current_q[d] = q[d];\n }\n\n // Sample p ~ N(0,1) - use simple LCG for speed, or curand\n curandState state;\n curand_init(1234ULL + b, 0, 0, &state);\n float sum_p2 = 0.0f;\n for (int d = 0; d < dim; d++) {\n float pd = curand_normal(&state);\n p[d] = pd;\n current_p[d] = pd;\n sum_p2 += pd * pd;\n }\n\n // Half step momentum: p += -0.5 * eps * q\n float half_eps = 0.5f * eps;\n for (int d = 0; d < dim; d++) {\n p[d] -= half_eps * q[d];\n }\n\n // Leapfrog steps\n for (int step = 0; step < num_leapfrog_steps - 1; step++) {\n // q += eps * p\n for (int d = 0; d < dim; d++) {\n q[d] += eps * p[d];\n }\n // p += eps * (-q)\n for (int d = 0; d < dim; d++) {\n p[d] -= eps * q[d];\n }\n }\n\n // Final full step\n for (int d = 0; d < dim; d++) {\n q[d] += eps * p[d];\n }\n for (int d = 0; d < dim; d++) {\n p[d] -= half_eps * q[d];\n }\n\n // p = -p\n float sum_p2_new = 0.0f;\n for (int d = 0; d < dim; d++) {\n p[d] = -p[d];\n sum_p2_new += p[d] * p[d];\n }\n\n // Compute energies\n float current_H = 0.0f;\n float proposed_H = 0.0f;\n for (int d = 0; d < dim; d++) {\n current_H += 0.5f * current_q[d] * current_q[d];\n proposed_H += 0.5f * q[d] * q[d];\n }\n current_H += 0.5f * sum_p2;\n proposed_H += 0.5f * sum_p2_new;\n\n float log_alpha = current_H - proposed_H;\n\n // Random\n float log_u = logf(curand_uniform(&state) + 1e-10f);\n bool acc = log_u < log_alpha;\n\n // Write output\n for (int d = 0; d < dim; d++) {\n q_out[b * dim + d] = acc ? q[d] : current_q[d];\n }\n accepted_out[b] = acc;\n}\n\nstd::vector<torch::Tensor> hmc_forward_full(\n torch::Tensor q,\n float step_size,\n int num_leapfrog_steps,\n int dim\n) {\n const int batch_size = q.size(0);\n auto q_out = torch::empty_like(q);\n auto accepted = torch::empty({batch_size}, q.options().dtype(torch::kBool));\n\n const int threads = 128;\n const int blocks = (batch_size + threads - 1) / threads;\n size_t shared_mem_size = threads * 4 * dim * sizeof(float);\n\n hmc_leapfrog_kernel<<<blocks, threads, shared_mem_size>>>(\n q.data_ptr<float>(),\n q_out.data_ptr<float>(),\n accepted.data_ptr<bool>(),\n batch_size,\n dim,\n step_size,\n num_leapfrog_steps\n );\n\n return {q_out, accepted};\n}\n'''\n\ncpp_source = r'''\n#include <torch/extension.h>\nstd::vector<torch::Tensor> hmc_forward_full(torch::Tensor q, float step_size, int num_leapfrog_steps, int dim);\n'''\n\nmodule = load_inline(\n name='hmc_cuda',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['hmc_forward_full'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, dim, num_leapfrog_steps=10, step_size=0.1):\n super().__init__()\n self.dim = dim\n self.num_leapfrog_steps = num_leapfrog_steps\n self.step_size = step_size\n self.register_buffer(\"mu\", torch.zeros(dim))\n\n def forward(self, q: torch.Tensor) -> tuple:\n results = module.hmc_forward_full(q.contiguous(), self.step_size, self.num_leapfrog_steps, self.dim)\n return results[0], results[1]\n```"
}
result ← submit_kernel (fail)
failsubmit_kernel
submit_kernel FAILED: kernel did not compile.
Syntax error in custom generated code or ModelNew not found