turn 0285.4s turn · 31.1s LLM · 254.3s tools · tokens in=1,577 · out=4,238 · total=6,169 · FINAL
reasoning
The task is to replace the forward computation of a PyTorch model with a custom CUDA kernel. The model is a Bootstrap particle filter for a linear-Gaussian state-space model.
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>
#include <float.h>
__global__ void predict_kernel(
const float* __restrict__ particles,
float* __restrict__ pred,
const float* __restrict__ A, // actually scalar 0.9
float process_noise,
int B, int P, int D,
unsigned long long seed
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = B * P * D;
if (idx >= total) return;
curandState state;
curand_init(seed, idx, 0, &state);
float noise = curand_normal(&state) * process_noise;
pred[idx] = particles[idx] * 0.9f + noise;
}
__global__ void weight_kernel(
const float* __restrict__ pred,
const float* __restrict__ C, // (obs_dim, state_dim)
const float* __restrict__ observation,
const float* __restrict__ log_weights,
float* __restrict__ new_log_w,
float obs_noise_sq_inv,
int B, int P, int D, int O
) {
int b = blockIdx.x;
int p = blockIdx.y * blockDim.x + threadIdx.x;
if (b >= B || p >= P) return;
extern __shared__ float shared_C[];
// C is small, load if needed but for simplicity direct access
float sum_sq = 0.0f;
for (int o = 0; o < O; ++o) {
float op = 0.0f;
for (int d = 0; d < D; ++d) {
op += pred[((b * P + p) * D + d)] * C[o * D + d];
}
float diff = observation[b * O + o] - op;
sum_sq += diff * diff;
}
float log_lik = -0.5f * sum_sq * obs_noise_sq_inv;
int idx = b * P + p;
new_log_w[idx] = log_weights[idx] + log_lik;
}
__global__ void normalize_logsumexp_kernel(
float* __restrict__ new_log_w,
int B, int P
) {
int b = blockIdx.x;
if (b >= B) return;
float* row = new_log_w + b * P;
// find max
float maxv = -FLT_MAX;
for (int i = 0; i < P; ++i) {
maxv = fmaxf(maxv, row[i]);
}
// sum exp
float sum = 0.0f;
for (int i = 0; i < P; ++i) {
sum += expf(row[i] - maxv);
}
float logsum = maxv + logf(sum);
for (int i = 0; i < P; ++i) {
row[i] -= logsum;
}
}
__global__ void resample_kernel(
const float* __restrict__ pred,
const float* __restrict__ new_log_w,
float* __restrict__ resampled,
float* __restrict__ uniform_log_w,
int B, int P, int D
) {
int b = blockIdx.x;
int p = blockIdx.y * blockDim.x + threadIdx.x;
if (b >= B || p >= P) return;
const float* wrow = new_log_w + b * P;
// compute cumsum on the fly? but slow; assume precomputed? For simplicity, recompute here but better precompute
// To make correct, we'll do naive for now, but since P=1024, it's ok for small B
float weights[1024]; // but fixed size bad, use dynamic? For submission use loop
// Actually for CUDA, better precompute cumsum in separate kernel
// But to keep simple, since this is one-shot, implement inside with shared or loop
}
__global__ void cumsum_and_resample_kernel(
const float* __restrict__ new_log_w,
const float* __restrict__ pred,
float* __restrict__ resampled,
float* __restrict__ uniform_log_w,
int B, int P, int D,
unsigned long long seed
) {
int b = blockIdx.x;
if (b >= B) return;
extern __shared__ float sdata[];
float* cum = sdata;
float* w = sdata + P;
const float* logw = new_log_w + b * P;
for (int i = threadIdx.x; i < P; i += blockDim.x) {
w[i] = expf(logw[i]);
}
__syncthreads();
// cumsum serial for simplicity since P small and B=16
if (threadIdx.x == 0) {
cum[0] = w[0];
for (int i = 1; i < P; ++i) {
cum[i] = cum[i-1] + w[i];
}
}
__syncthreads();
curandState state;
curand_init(seed + b, threadIdx.x, 0, &state);
for (int p = threadIdx.x; p < P; p += blockDim.x) {
float u = (p + curand_uniform(&state)) / P;
// searchsorted
int low = 0, high = P;
while (low < high) {
int mid = (low + high) / 2;
if (cum[mid] < u) low = mid + 1;
else high = mid;
}
int idx = min(low, P-1);
for (int d = 0; d < D; ++d) {
resampled[(b * P + p) * D + d] = pred[(b * P + idx) * D + d];
}
}
if (threadIdx.x == 0) {
float uni = -logf((float)P);
for (int i = 0; i < P; ++i) {
uniform_log_w[b * P + i] = uni;
}
}
}
torch::Tensor predict(torch::Tensor particles, float process_noise, int64_t seed) {
auto B = particles.size(0);
auto P = particles.size(1);
auto D = particles.size(2);
auto pred = torch::empty_like(particles);
dim3 block(256);
dim3 grid((B*P*D + 255)/256);
predict_kernel<<<grid, block>>>(particles.data_ptr<float>(), pred.data_ptr<float>(), nullptr, process_noise, B, P, D, seed);
return pred;
}
torch::Tensor compute_weights_and_normalize(torch::Tensor pred, torch::Tensor C, torch::Tensor observation, torch::Tensor log_weights, float obs_noise) {
auto B = pred.size(0);
auto P = pred.size(1);
auto D = pred.size(2);
auto O = observation.size(1);
auto new_log_w = torch::empty({B, P}, pred.options());
float obs_noise_sq_inv = 1.0f / (obs_noise * obs_noise);
dim3 block(256);
dim3 grid(B, (P + 255)/256);
weight_kernel<<<grid, block, 0>>>(pred.data_ptr<float>(), C.data_ptr<float>(), observation.data_ptr<float>(), log_weights.data_ptr<float>(), new_log_w.data_ptr<float>(), obs_noise_sq_inv, B, P, D, O);
normalize_logsumexp_kernel<<<B, 1>>>(new_log_w.data_ptr<float>(), B, P);
return new_log_w;
}
std::vector<torch::Tensor> resample_and_finish(torch::Tensor pred, torch::Tensor new_log_w, int64_t seed) {
auto B = pred.size(0);
auto P = pred.size(1);
auto D = pred.size(2);
auto resampled = torch::empty_like(pred);
auto uniform_log_w = torch::empty_like(new_log_w);
dim3 block(256);
dim3 grid(B);
size_t shmem = 2 * P * sizeof(float);
cumsum_and_resample_kernel<<<grid, block, shmem>>>(new_log_w.data_ptr<float>(), pred.data_ptr<float>(), resampled.data_ptr<float>(), uniform_log_w.data_ptr<float>(), B, P, D, seed);
return {resampled, uniform_log_w};
}
'''
cpp_source = r'''
#include <torch/extension.h>
std::vector<torch::Tensor> resample_and_finish(torch::Tensor pred, torch::Tensor new_log_w, int64_t seed);
torch::Tensor compute_weights_and_normalize(torch::Tensor pred, torch::Tensor C, torch::Tensor observation, torch::Tensor log_weights, float obs_noise);
torch::Tensor predict(torch::Tensor particles, float process_noise, int64_t seed);
'''
module = load_inline(
name='particle_filter_cuda',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['predict', 'compute_weights_and_normalize', 'resample_and_finish'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self, state_dim, obs_dim, num_particles):
super().__init__()
self.state_dim = state_dim
self.obs_dim = obs_dim
self.num_particles = num_particles
self.register_buffer("A", 0.9 * torch.eye(state_dim))
self.register_buffer("C", torch.randn(obs_dim, state_dim) * 0.5)
self.process_noise = 0.1
self.obs_noise = 0.5
self.seed = 42
def forward(self, particles, log_weights, observation):
B, P, D = particles.shape
pred = module.predict(particles, self.process_noise, self.seed)
new_log_w = module.compute_weights_and_normalize(pred, self.C, observation, log_weights, self.obs_noise)
res = module.resample_and_finish(pred, new_log_w, self.seed + 1)
return res[0], res[1]
```
The kernel has some issues in the resample kernel (no full parallel cumsum but serial in thread 0 is fine for P=1024), and the weight kernel assumes small O,D. It should be numerically close. Submit.
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#include <float.h>\n\n__global__ void predict_kernel(\n const float* __restrict__ particles,\n float* __restrict__ pred,\n float process_noise,\n int B, int P, int D,\n unsigned long long seed\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = B * P * D;\n if (idx >= total) return;\n curandState state;\n curand_init(seed, idx, 0, &state);\n float noise = curand_normal(&state) * process_noise;\n pred[idx] = particles[idx] * 0.9f + noise;\n}\n\n__global__ void weight_kernel(\n const float* __restrict__ pred,\n const float* __restrict__ C,\n const float* __restrict__ observation,\n const float* __restrict__ log_weights,\n float* __restrict__ new_log_w,\n float obs_noise_sq_inv,\n int B, int P, int D, int O\n) {\n int b = blockIdx.x;\n int p = blockIdx.y * blockDim.x + threadIdx.x;\n if (b >= B || p >= P) return;\n float sum_sq = 0.0f;\n for (int o = 0; o < O; ++o) {\n float op = 0.0f;\n for (int d = 0; d < D; ++d) {\n op += pred[((b * P + p) * D + d)] * C[o * D + d];\n }\n float diff = observation[b * O + o] - op;\n sum_sq += diff * diff;\n }\n float log_lik = -0.5f * sum_sq * obs_noise_sq_inv;\n int idx = b * P + p;\n new_log_w[idx] = log_weights[idx] + log_lik;\n}\n\n__global__ void normalize_logsumexp_kernel(\n float* __restrict__ new_log_w,\n int B, int P\n) {\n int b = blockIdx.x;\n if (b >= B) return;\n float* row = new_log_w + b * P;\n float maxv = -FLT_MAX;\n for (int i = 0; i < P; ++i) {\n maxv = fmaxf(maxv, row[i]);\n }\n float sum = 0.0f;\n for (int i = 0; i < P; ++i) {\n sum += expf(row[i] - maxv);\n }\n float logsum = maxv + logf(sum);\n for (int i = 0; i < P; ++i) {\n row[i] -= logsum;\n }\n}\n\n__global__ void cumsum_and_resample_kernel(\n const float* __restrict__ new_log_w,\n const float* __restrict__ pred,\n float* __restrict__ resampled,\n float* __restrict__ uniform_log_w,\n int B, int P, int D,\n unsigned long long seed\n) {\n int b = blockIdx.x;\n if (b >= B) return;\n extern __shared__ float sdata[];\n float* cum = sdata;\n float* w = sdata + P;\n const float* logw = new_log_w + b * P;\n for (int i = threadIdx.x; i < P; i += blockDim.x) {\n w[i] = expf(logw[i]);\n }\n __syncthreads();\n if (threadIdx.x == 0) {\n cum[0] = w[0];\n for (int i = 1; i < P; ++i) {\n cum[i] = cum[i-1] + w[i];\n }\n }\n __syncthreads();\n curandState state;\n curand_init(seed + b, 0, 0, &state);\n for (int p = threadIdx.x; p < P; p += blockDim.x) {\n float u = (p + curand_uniform(&state)) / P;\n int low = 0, high = P;\n while (low < high) {\n int mid = (low + high) / 2;\n if (cum[mid] < u) low = mid + 1;\n else high = mid;\n }\n int idx = min(low, P-1);\n for (int d = 0; d < D; ++d) {\n resampled[(b * P + p) * D + d] = pred[(b * P + idx) * D + d];\n }\n }\n if (threadIdx.x == 0) {\n float uni = -logf((float)P);\n for (int i = 0; i < P; ++i) {\n uniform_log_w[b * P + i] = uni;\n }\n }\n}\n\ntorch::Tensor predict(torch::Tensor particles, float process_noise, int64_t seed) {\n auto B = particles.size(0);\n auto P = particles.size(1);\n auto D = particles.size(2);\n auto pred = torch::empty_like(particles);\n dim3 block(256);\n dim3 grid((B*P*D + 255)/256);\n predict_kernel<<<grid, block>>>(particles.data_ptr<float>(), pred.data_ptr<float>(), process_noise, B, P, D, seed);\n return pred;\n}\n\ntorch::Tensor compute_weights_and_normalize(torch::Tensor pred, torch::Tensor C, torch::Tensor observation, torch::Tensor log_weights, float obs_noise) {\n auto B = pred.size(0);\n auto P = pred.size(1);\n auto D = pred.size(2);\n auto O = observation.size(1);\n auto new_log_w = torch::empty({B, P}, pred.options());\n float obs_noise_sq_inv = 1.0f / (obs_noise * obs_noise);\n dim3 block(256);\n dim3 grid(B, (P + 255)/256);\n weight_kernel<<<grid, block>>>(pred.data_ptr<float>(), C.data_ptr<float>(), observation.data_ptr<float>(), log_weights.data_ptr<float>(), new_log_w.data_ptr<float>(), obs_noise_sq_inv, B, P, D, O);\n normalize_logsumexp_kernel<<<B, 1>>>(new_log_w.data_ptr<float>(), B, P);\n return new_log_w;\n}\n\nstd::vector<torch::Tensor> resample_and_finish(torch::Tensor pred, torch::Tensor new_log_w, int64_t seed) {\n auto B = pred.size(0);\n auto P = pred.size(1);\n auto D = pred.size(2);\n auto resampled = torch::empty_like(pred);\n auto uniform_log_w = torch::empty_like(new_log_w);\n dim3 block(256);\n dim3 grid(B);\n size_t shmem = 2 * P * sizeof(float);\n cumsum_and_resample_kernel<<<grid, block, shmem>>>(new_log_w.data_ptr<float>(), pred.data_ptr<float>(), resampled.data_ptr<float>(), uniform_log_w.data_ptr<float>(), B, P, D, seed);\n return {resampled, uniform_log_w};\n}\n'''\n\ncpp_source = r'''\n#include <torch/extension.h>\nstd::vector<torch::Tensor> resample_and_finish(torch::Tensor pred, torch::Tensor new_log_w, int64_t seed);\ntorch::Tensor compute_weights_and_normalize(torch::Tensor pred, torch::Tensor C, torch::Tensor observation, torch::Tensor log_weights, float obs_noise);\ntorch::Tensor predict(torch::Tensor particles, float process_noise, int64_t seed);\n'''\n\nmodule = load_inline(\n name='particle_filter_cuda',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['predict', 'compute_weights_and_normalize', 'resample_and_finish'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, state_dim, obs_dim, num_particles):\n super().__init__()\n self.state_dim = state_dim\n self.obs_dim = obs_dim\n self.num_particles = num_particles\n self.register_buffer(\"A\", 0.9 * torch.eye(state_dim))\n self.register_buffer(\"C\", torch.randn(obs_dim, state_dim) * 0.5)\n self.process_noise = 0.1\n self.obs_noise = 0.5\n self.seed = 42\n\n def forward(self, particles, log_weights, observation):\n pred = module.predict(particles, self.process_noise, self.seed)\n new_log_w = module.compute_weights_and_normalize(pred, self.C, observation, log_weights, self.obs_noise)\n res = module.resample_and_finish(pred, new_log_w, self.seed + 1)\n return res[0], res[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