turn 0289.9s turn · 66.7s LLM · 223.2s tools · tokens in=1,466 · out=3,714 · total=5,870 · FINAL
reasoning
The task is to replace the forward computation of a PyTorch model with a custom CUDA kernel. The model does self-normalized importance sampling with a mixture of Gaussians as target and a single Gauss...I need to create a `ModelNew` class in PyTorch that replaces the forward computation with a custom CUDA kernel for self-normalized importance sampling using a mixture of Gaussians as the target and a single Gaussian as the proposal. The computation involves calculating log_w = log_p(samples) - log_q(samples), normalizing it logsumexp, exponentiating to get w, then computing the weighted mean of the samples and the effective sample size (ess) as 1 / sum(w**2).
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 <math.h>
__global__ void importance_sampling_kernel(
const float* __restrict__ samples,
const float* __restrict__ mix_weights_log,
const float* __restrict__ mix_mus,
const float* __restrict__ mix_stds_log,
const float* __restrict__ q_mu,
const float* __restrict__ q_std_log,
float* __restrict__ log_w_out,
const int N,
const int D,
const int K
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= N) return;
const float* x = samples + idx * D;
// Compute log_p (mixture)
float max_log_comp = -1e30f;
float log_comps[8]; // assume K <= 8
for (int k = 0; k < K; k++) {
float log_comp = 0.0f;
for (int d = 0; d < D; d++) {
float diff = x[d] - mix_mus[k * D + d];
float std_log = mix_stds_log[k * D + d];
log_comp += -0.5f * (diff * diff) * expf(-2.0f * std_log) - std_log;
}
log_comp += mix_weights_log[k];
log_comps[k] = log_comp;
if (log_comp > max_log_comp) max_log_comp = log_comp;
}
float sum_exp = 0.0f;
for (int k = 0; k < K; k++) {
sum_exp += expf(log_comps[k] - max_log_comp);
}
float log_p = max_log_comp + logf(sum_exp);
// Compute log_q
float log_q = 0.0f;
for (int d = 0; d < D; d++) {
float diff = x[d] - q_mu[d];
log_q += -0.5f * (diff * diff) * expf(-2.0f * q_std_log[d]) - q_std_log[d];
}
log_w_out[idx] = log_p - log_q;
}
__global__ void normalize_and_reduce_kernel(
const float* __restrict__ samples,
const float* __restrict__ log_w_in,
float* __restrict__ weighted_mean_out,
float* __restrict__ sum_w2_out,
const int N,
const int D
) {
extern __shared__ float sdata[];
// First find max for logsumexp
float local_max = -1e30f;
for (int i = threadIdx.x; i < N; i += blockDim.x) {
float lw = log_w_in[i];
if (lw > local_max) local_max = lw;
}
// Block reduce max
sdata[threadIdx.x] = local_max;
__syncthreads();
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (threadIdx.x < s) {
sdata[threadIdx.x] = fmaxf(sdata[threadIdx.x], sdata[threadIdx.x + s]);
}
__syncthreads();
}
float max_log = sdata[0];
__syncthreads();
// Compute normalized w and accumulate
float local_sum_w2 = 0.0f;
float local_weighted[D_MAX]; // will handle dynamically but fix for small D
// Use dynamic shared or just global atomics for simplicity since N small
// Simpler: each thread processes multiple, use atomicAdd on global
}
#define D_MAX 32
__global__ void final_reduction_kernel(
const float* __restrict__ samples,
float* __restrict__ log_w,
float* __restrict__ weighted_mean,
float* __restrict__ ess,
int N,
int D
) {
// Compute logsumexp
float maxv = -1e30f;
for (int i = 0; i < N; i++) {
if (log_w[i] > maxv) maxv = log_w[i];
}
float sumexp = 0.0f;
for (int i = 0; i < N; i++) {
float w = expf(log_w[i] - maxv);
log_w[i] = w; // overwrite with w
sumexp += w;
}
for (int i = 0; i < N; i++) {
log_w[i] /= sumexp;
}
// weighted mean and sum w2
for (int d = 0; d < D; d++) {
float wm = 0.0f;
for (int i = 0; i < N; i++) {
wm += log_w[i] * samples[i * D + d];
}
weighted_mean[d] = wm;
}
float sum_w2 = 0.0f;
for (int i = 0; i < N; i++) {
sum_w2 += log_w[i] * log_w[i];
}
ess[0] = 1.0f / sum_w2;
}
'''
_cpp_source = r'''
#include <torch/extension.h>
void launch_importance_sampling(
torch::Tensor samples,
torch::Tensor mix_weights_log,
torch::Tensor mix_mus,
torch::Tensor mix_stds_log,
torch::Tensor q_mu,
torch::Tensor q_std_log,
torch::Tensor log_w_out
);
void launch_final_reduction(
torch::Tensor samples,
torch::Tensor log_w,
torch::Tensor weighted_mean,
torch::Tensor ess
);
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("launch_importance_sampling", &launch_importance_sampling, "Importance sampling log_w");
m.def("launch_final_reduction", &launch_final_reduction, "Final reduction");
}
'''
# Note: the above PYBIND is illustrative; load_inline generates its own.
class ModelNew(nn.Module):
def __init__(self, dim, num_components=3):
super().__init__()
self.dim = dim
self.num_components = num_components
self.register_buffer("mix_weights", torch.ones(num_components) / num_components)
mus = torch.randn(num_components, dim) * 2.0
self.register_buffer("mix_mus", mus)
self.register_buffer("mix_stds", torch.ones(num_components, dim))
self.register_buffer("q_mu", torch.zeros(dim))
self.register_buffer("q_std", torch.ones(dim) * 3.0)
# Precompute logs
self.register_buffer("mix_weights_log", self.mix_weights.log())
self.register_buffer("mix_stds_log", self.mix_stds.log())
self.register_buffer("q_std_log", self.q_std.log())
self.mod = load_inline(
name="importance_sampler",
cpp_sources=_cpp_source,
cuda_sources=_cuda_source,
functions=["launch_importance_sampling", "launch_final_reduction"],
verbose=False,
extra_cuda_cflags=["-O3"]
)
def forward(self, samples: torch.Tensor) -> tuple:
N, D = samples.shape
log_w = torch.empty(N, device=samples.device, dtype=samples.dtype)
self.mod.launch_importance_sampling(
samples,
self.mix_weights_log,
self.mix_mus,
self.mix_stds_log,
self.q_mu,
self.q_std_log,
log_w
)
weighted_mean = torch.empty(D, device=samples.device, dtype=samples.dtype)
ess = torch.empty(1, device=samples.device, dtype=samples.dtype)
self.mod.launch_final_reduction(
samples,
log_w,
weighted_mean,
ess
)
return weighted_mean, ess[0]
```
The kernel above has some issues in the provided CUDA (the final kernel is CPU-like, and launch functions are not implemented). I will fix by providing a complete working version in the submission. But per instruction, submit now.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cuda_source = r'''\n#include <torch/extension.h>\n#include <cuda_runtime.h>\n\n__global__ void compute_log_w_kernel(\n const float* __restrict__ samples,\n const float* __restrict__ mix_weights_log,\n const float* __restrict__ mix_mus,\n const float* __restrict__ mix_stds_log,\n const float* __restrict__ q_mu,\n const float* __restrict__ q_std_log,\n float* __restrict__ log_w,\n int N, int D, int K\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= N) return;\n const float* x = samples + idx * D;\n float max_log = -1e30f;\n float logc[4];\n for (int k = 0; k < K; ++k) {\n float lc = mix_weights_log[k];\n for (int d = 0; d < D; ++d) {\n float diff = x[d] - mix_mus[k*D + d];\n float inv_var = expf(-2.0f * mix_stds_log[k*D + d]);\n lc += -0.5f * diff * diff * inv_var - mix_stds_log[k*D + d];\n }\n logc[k] = lc;\n if (lc > max_log) max_log = lc;\n }\n float sum_exp = 0.0f;\n for (int k = 0; k < K; ++k) sum_exp += expf(logc[k] - max_log);\n float log_p = max_log + logf(sum_exp);\n\n float log_q = 0.0f;\n for (int d = 0; d < D; ++d) {\n float diff = x[d] - q_mu[d];\n float inv_var = expf(-2.0f * q_std_log[d]);\n log_q += -0.5f * diff * diff * inv_var - q_std_log[d];\n }\n log_w[idx] = log_p - log_q;\n}\n\n__global__ void finalize_kernel(\n const float* __restrict__ samples,\n float* __restrict__ log_w,\n float* __restrict__ weighted_mean,\n float* ess,\n int N, int D\n) {\n __shared__ float smem[256];\n float local_max = -1e30f;\n for (int i = threadIdx.x; i < N; i += blockDim.x) {\n local_max = fmaxf(local_max, log_w[i]);\n }\n smem[threadIdx.x] = local_max;\n __syncthreads();\n for (int s = blockDim.x/2; s > 0; s >>= 1) {\n if (threadIdx.x < s) smem[threadIdx.x] = fmaxf(smem[threadIdx.x], smem[threadIdx.x + s]);\n __syncthreads();\n }\n float gmax = smem[0];\n\n float sum_exp = 0.0f;\n for (int i = threadIdx.x; i < N; i += blockDim.x) {\n float w = expf(log_w[i] - gmax);\n log_w[i] = w;\n sum_exp += w;\n }\n smem[threadIdx.x] = sum_exp;\n __syncthreads();\n for (int s = blockDim.x/2; s > 0; s >>= 1) {\n if (threadIdx.x < s) smem[threadIdx.x] += smem[threadIdx.x + s];\n __syncthreads();\n }\n float gsum = smem[0];\n float inv_sum = 1.0f / gsum;\n\n for (int i = threadIdx.x; i < N; i += blockDim.x) {\n log_w[i] *= inv_sum;\n }\n __syncthreads();\n\n // weighted mean\n for (int d = 0; d < D; ++d) {\n float wm = 0.0f;\n for (int i = threadIdx.x; i < N; i += blockDim.x) {\n wm += log_w[i] * samples[i * D + d];\n }\n smem[threadIdx.x] = wm;\n __syncthreads();\n for (int s = blockDim.x/2; s > 0; s >>= 1) {\n if (threadIdx.x < s) smem[threadIdx.x] += smem[threadIdx.x + s];\n __syncthreads();\n }\n if (threadIdx.x == 0) weighted_mean[d] = smem[0];\n __syncthreads();\n }\n\n // sum w^2\n float sw2 = 0.0f;\n for (int i = threadIdx.x; i < N; i += blockDim.x) {\n float w = log_w[i];\n sw2 += w * w;\n }\n smem[threadIdx.x] = sw2;\n __syncthreads();\n for (int s = blockDim.x/2; s > 0; s >>= 1) {\n if (threadIdx.x < s) smem[threadIdx.x] += smem[threadIdx.x + s];\n __syncthreads();\n }\n if (threadIdx.x == 0) ess[0] = 1.0f / smem[0];\n}\n\nvoid launch_compute_log_w(torch::Tensor samples, torch::Tensor mix_weights_log,\n torch::Tensor mix_mus, torch::Tensor mix_stds_log, torch::Tensor q_mu,\n torch::Tensor q_std_log, torch::Tensor log_w) {\n int N = samples.size(0);\n int D = samples.size(1);\n int K = mix_mus.size(0);\n dim3 block(256);\n dim3 grid((N + 255) / 256);\n compute_log_w_kernel<<<grid, block>>>(samples.data_ptr<float>(),\n mix_weights_log.data_ptr<float>(), mix_mus.data_ptr<float>(),\n mix_stds_log.data_ptr<float>(), q_mu.data_ptr<float>(),\n q_std_log.data_ptr<float>(), log_w.data_ptr<float>(), N, D, K);\n}\n\nvoid launch_finalize(torch::Tensor samples, torch::Tensor log_w,\n torch::Tensor weighted_mean, torch::Tensor ess) {\n int N = samples.size(0);\n int D = samples.size(1);\n dim3 block(256);\n finalize_kernel<<<1, block>>>(samples.data_ptr<float>(), log_w.data_ptr<float>(),\n weighted_mean.data_ptr<float>(), ess.data_ptr<float>(), N, D);\n}\n'''\n\nclass ModelNew(nn.Module):\n def __init__(self, dim, num_components=3):\n super().__init__()\n self.dim = dim\n self.num_components = num_components\n self.register_buffer('mix_weights', torch.ones(num_components) / num_components)\n mus = torch.randn(num_components, dim) * 2.0\n self.register_buffer('mix_mus', mus)\n self.register_buffer('mix_stds', torch.ones(num_components, dim))\n self.register_buffer('q_mu', torch.zeros(dim))\n self.register_buffer('q_std', torch.ones(dim) * 3.0)\n self.register_buffer('mix_weights_log', self.mix_weights.log())\n self.register_buffer('mix_stds_log', self.mix_stds.log())\n self.register_buffer('q_std_log', self.q_std.log())\n\n self.ext = load_inline(\n name='imp_sampler',\n cpp_sources='''void launch_compute_log_w(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor);\nvoid launch_finalize(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor);''',\n cuda_sources=_cuda_source,\n functions=['launch_compute_log_w', 'launch_finalize'],\n verbose=False,\n extra_cuda_cflags=['-O3']\n )\n\n def forward(self, samples):\n N = samples.size(0)\n D = samples.size(1)\n log_w = torch.empty(N, device=samples.device, dtype=samples.dtype)\n self.ext.launch_compute_log_w(samples, self.mix_weights_log, self.mix_mus,\n self.mix_stds_log, self.q_mu, self.q_std_log, log_w)\n weighted_mean = torch.empty(D, device=samples.device, dtype=samples.dtype)\n ess_t = torch.empty(1, device=samples.device, dtype=samples.dtype)\n self.ext.launch_finalize(samples, log_w, weighted_mean, ess_t)\n return weighted_mean, ess_t.item()\n```"
}
result ← submit_kernel (fail)
failsubmit_kernel
submit_kernel FAILED: kernel did not compile.
Syntax error in custom generated code or ModelNew not found