call → compile_kernel
callcompile_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 <cmath>\n\nconstexpr int BLOCK_SIZE = 256;\nconstexpr int WARP_SIZE = 32;\n\n// Helper for warp-level reduction\n__inline__ __device__ float warp_reduce_max(float val) {\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n val = max(val, __shfl_down_sync(0xffffffff, val, offset));\n }\n return val;\n}\n\n__inline__ __device__ float warp_reduce_sum(float val) {\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n val += __shfl_down_sync(0xffffffff, val, offset);\n }\n return val;\n}\n\n// Device function to compute log_p for a single sample\n__device__ float compute_log_p_single(\n const float* __restrict__ sample,\n const float* __restrict__ mix_mus,\n const float* __restrict__ mix_stds,\n const float* __restrict__ mix_weights,\n int dim, int num_components) {\n \n float max_log_comp = -1e30f;\n float sum_exp = 0.0f;\n \n // Compute log-sum-exp across mixture components\n for (int k = 0; k < num_components; k++) {\n const float* mu_k = mix_mus + k * dim;\n const float* std_k = mix_stds + k * dim;\n float log_weight = logf(mix_weights[k]);\n \n float log_comp = log_weight;\n for (int d = 0; d < dim; d++) {\n float diff = sample[d] - mu_k[d];\n float std = std_k[d];\n log_comp += -0.5f * (diff * diff) / (std * std) - logf(std);\n }\n \n if (k == 0 || log_comp > max_log_comp) {\n max_log_comp = log_comp;\n }\n sum_exp += expf(log_comp - max_log_comp); // Will renormalize later\n }\n \n // Actually need proper logsumexp\n // Recompute with correct max\n sum_exp = 0.0f;\n for (int k = 0; k < num_components; k++) {\n const float* mu_k = mix_mus + k * dim;\n const float* std_k = mix_stds + k * dim;\n float log_weight = logf(mix_weights[k]);\n \n float log_comp = log_weight;\n for (int d = 0; d < dim; d++) {\n float diff = sample[d] - mu_k[d];\n float std = std_k[d];\n log_comp += -0.5f * (diff * diff) / (std * std) - logf(std);\n }\n \n sum_exp += expf(log_comp - max_log_comp);\n }\n \n return max_log_comp + logf(sum_exp);\n}\n\n// Device function to compute log_q for a single sample\n__device__ float compute_log_q_single(\n const float* __restrict__ sample,\n const float* __restrict__ q_mu,\n const float* __restrict__ q_std,\n int dim) {\n \n float log_q = 0.0f;\n for (int d = 0; d < dim; d++) {\n float diff = sample[d] - q_mu[d];\n float std = q_std[d];\n log_q += -0.5f * (diff * diff) / (std * std) - logf(std);\n }\n return log_q;\n}\n\n// Kernel 1: Compute log weights (log_p - log_q) for all samples\n__global__ void compute_log_weights_kernel(\n const float* __restrict__ samples,\n const float* __restrict__ mix_mus,\n const float* __restrict__ mix_stds,\n const float* __restrict__ mix_weights,\n const float* __restrict__ q_mu,\n const float* __restrict__ q_std,\n float* __restrict__ log_w,\n int N, int dim, int num_components) {\n \n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < N) {\n const float* sample = samples + idx * dim;\n float log_p = compute_log_p_single(sample, mix_mus, mix_stds, mix_weights, dim, num_components);\n float log_q = compute_log_q_single(sample, q_mu, q_std, dim);\n log_w[idx] = log_p - log_q;\n }\n}\n\n// Kernel 2: Compute logsumexp reduction\n// Uses two-stage reduction: first block-level, then global\n__global__ void logsumexp_kernel(\n const float* __restrict__ log_w,\n float* __restrict__ block_maxes,\n float* __restrict__ block_sums,\n int N) {\n \n __shared__ float s_max[BLOCK_SIZE / WARP_SIZE];\n __shared__ float s_sum[BLOCK_SIZE / WARP_SIZE];\n \n int tid = threadIdx.x;\n int warp_id = tid / WARP_SIZE;\n int lane_id = tid % WARP_SIZE;\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n \n // Initialize with thread's value or very negative number\n float my_val = (idx < N) ? log_w[idx] : -1e30f;\n \n // Warp-level max reduction\n float warp_max = warp_reduce_max(my_val);\n \n // Lane 0 of each warp writes to shared memory\n if (lane_id == 0) {\n s_max[warp_id] = warp_max;\n }\n __syncthreads();\n \n // First warp reduces the warp maxes\n float block_max;\n if (warp_id == 0) {\n block_max = (lane_id < (BLOCK_SIZE / WARP_SIZE)) ? s_max[lane_id] : -1e30f;\n block_max = warp_reduce_max(block_max);\n if (lane_id == 0) {\n block_maxes[blockIdx.x] = block_max;\n }\n }\n __syncthreads();\n \n // Get block max from shared memory (broadcast from thread 0)\n block_max = s_max[0];\n \n // Compute exp(log_w - block_max) for this thread\n float my_exp = (idx < N) ? expf(my_val - block_max) : 0.0f;\n \n // Warp-level sum reduction\n float warp_sum = warp_reduce_sum(my_exp);\n \n // Lane 0 of each warp writes to shared memory\n if (lane_id == 0) {\n s_sum[warp_id] = warp_sum;\n }\n __syncthreads();\n \n // First warp reduces the warp sums\n float block_sum;\n if (warp_id == 0) {\n block_sum = (lane_id < (BLOCK_SIZE / WARP_SIZE)) ? s_sum[lane_id] : 0.0f;\n block_sum = warp_reduce_sum(block_sum);\n if (lane_id == 0) {\n block_sums[blockIdx.x] = block_sum;\n }\n }\n}\n\n// Kernel 3: Normalize weights and compute weighted mean and ESS\n__global__ void normalize_and_reduce_kernel(\n const float* __restrict__ samples,\n const float* __restrict__ log_w,\n float global_logsumexp,\n float* __restrict__ weighted_mean_partial,\n float* __restrict__ ess_partial,\n int N, int dim) {\n \n __shared__ float s_mean[BLOCK_SIZE][32]; // partial sums for mean (up to dim=32)\n __shared__ float s_ess[BLOCK_SIZE];\n \n int tid = threadIdx.x;\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n \n // Initialize shared memory\n for (int d = tid; d < dim && d < 32; d += blockDim.x) {\n s_mean[tid][d] = 0.0f;\n }\n s_ess[tid] = 0.0f;\n __syncthreads();\n \n if (idx < N) {\n float w = expf(log_w[idx] - global_logsumexp);\n const float* sample = samples + idx * dim;\n \n // Accumulate weighted mean and ESS\n for (int d = 0; d < dim && d < 32; d++) {\n atomicAdd(&s_mean[0][d], w * sample[d]);\n }\n atomicAdd(&s_ess[0], w * w);\n }\n __syncthreads();\n \n // Write partial results\n if (tid == 0) {\n for (int d = 0; d < dim && d < 32; d++) {\n weighted_mean_partial[blockIdx.x * dim + d] = s_mean[0][d];\n }\n ess_partial[blockIdx.x] = s_ess[0];\n }\n}\n\n// Final reduction kernel for weighted mean and ESS\n__global__ void final_reduce_kernel(\n const float* __restrict__ weighted_mean_partial,\n const float* __restrict__ ess_partial,\n float* __restrict__ weighted_mean_out,\n float* __restrict__ ess_out,\n int num_blocks, int dim) {\n \n __shared__ float s_mean[32];\n __shared__ float s_ess;\n \n int tid = threadIdx.x;\n \n if (tid < dim && tid < 32) {\n float sum = 0.0f;\n for (int b = 0; b < num_blocks; b++) {\n sum += weighted_mean_partial[b * dim + tid];\n }\n s_mean[tid] = sum;\n }\n \n if (tid == 0) {\n float sum = 0.0f;\n for (int b = 0; b < num_blocks; b++) {\n sum += ess_partial[b];\n }\n s_ess = sum;\n }\n __syncthreads();\n \n if (tid < dim && tid < 32) {\n weighted_mean_out[tid] = s_mean[tid];\n }\n if (tid == 0) {\n ess_out[0] = 1.0f / s_ess;\n }\n}\n\n// Alternative single-kernel approach for better performance\n__global__ void importance_sampling_kernel(\n const float* __restrict__ samples,\n const float* __restrict__ mix_mus,\n const float* __restrict__ mix_stds,\n const float* __restrict__ mix_weights,\n const float* __restrict__ q_mu,\n const float* __restrict__ q_std,\n float* __restrict__ log_w,\n float* __restrict__ block_maxes,\n float* __restrict__ block_sums,\n int N, int dim, int num_components) {\n \n extern __shared__ float shared_mem[];\n float* s_max = shared_mem; // BLOCK_SIZE / WARP_SIZE elements\n float* s_sum = s_max + (BLOCK_SIZE / WARP_SIZE); // BLOCK_SIZE / WARP_SIZE elements\n \n int tid = threadIdx.x;\n int warp_id = tid / WARP_SIZE;\n int lane_id = tid % WARP_SIZE;\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n \n // Stage 1: Compute log weights\n float my_log_w = -1e30f;\n if (idx < N) {\n const float* sample = samples + idx * dim;\n float log_p = compute_log_p_single(sample, mix_mus, mix_stds, mix_weights, dim, num_components);\n float log_q = compute_log_q_single(sample, q_mu, q_std, dim);\n my_log_w = log_p - log_q;\n log_w[idx] = my_log_w;\n }\n \n // Stage 2: Block-level logsumexp\n // Warp-level max reduction\n float warp_max = warp_reduce_max(my_log_w);\n \n if (lane_id == 0) {\n s_max[warp_id] = warp_max;\n }\n __syncthreads();\n \n float block_max;\n if (warp_id == 0) {\n block_max = (lane_id < (BLOCK_SIZE / WARP_SIZE)) ? s_max[lane_id] : -1e30f;\n block_max = warp_reduce_max(block_max);\n if (lane_id == 0) {\n block_maxes[blockIdx.x] = block_max;\n }\n }\n __syncthreads();\n \n // Broadcast block_max (read from shared or global depending on timing)\n __shared__ float s_block_max;\n if (tid == 0) {\n s_block_max = block_maxes[blockIdx.x];\n }\n __syncthreads();\n block_max = s_block_max;\n \n // Compute exp(log_w - block_max)\n float my_exp = (idx < N) ? expf(my_log_w - block_max) : 0.0f;\n \n // Warp-level sum reduction\n float warp_sum = warp_reduce_sum(my_exp);\n \n if (lane_id == 0) {\n s_sum[warp_id] = warp_sum;\n }\n __syncthreads();\n \n if (warp_id == 0) {\n float block_sum = (lane_id < (BLOCK_SIZE / WARP_SIZE)) ? s_sum[lane_id] : 0.0f;\n block_sum = warp_reduce_sum(block_sum);\n if (lane_id == 0) {\n block_sums[blockIdx.x] = block_sum;\n }\n }\n}\n\nstd::tuple<torch::Tensor, torch::Tensor> importance_sampling_cuda(\n torch::Tensor samples,\n torch::Tensor mix_mus,\n torch::Tensor mix_stds,\n torch::Tensor mix_weights,\n torch::Tensor q_mu,\n torch::Tensor q_std) {\n \n int N = samples.size(0);\n int dim = samples.size(1);\n int num_components = mix_mus.size(0);\n \n int num_blocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;\n \n // Allocate intermediate buffers\n auto options = samples.options();\n auto log_w = torch::empty({N}, options);\n auto block_maxes = torch::empty({num_blocks}, options);\n auto block_sums = torch::empty({num_blocks}, options);\n \n // Launch combined kernel\n size_t shared_mem_size = 2 * (BLOCK_SIZE / WARP_SIZE) * sizeof(float);\n \n importance_sampling_kernel<<<num_blocks, BLOCK_SIZE, shared_mem_size>>>(\n samples.data_ptr<float>(),\n mix_mus.data_ptr<float>(),\n mix_stds.data_ptr<float>(),\n mix_weights.data_ptr<float>(),\n q_mu.data_ptr<float>(),\n q_std.data_ptr<float>(),\n log_w.data_ptr<float>(),\n block_maxes.data_ptr<float>(),\n block_sums.data_ptr<float>(),\n N, dim, num_components\n );\n \n // Second pass: compute global logsumexp\n // We need to launch a small kernel or do it on CPU for the block results\n // For simplicity, use a small kernel\n auto block_maxes_2 = torch::empty({1}, options);\n auto block_sums_2 = torch::empty({1}, options);\n \n int num_blocks_2 = (num_blocks + BLOCK_SIZE - 1) / BLOCK_SIZE;\n \n // For now, do the final logsumexp reduction on CPU (small data)\n // This is not ideal but works for correctness\n torch::Tensor block_maxes_cpu = block_maxes.cpu();\n torch::Tensor block_sums_cpu = block_sums.cpu();\n \n float global_max = block_maxes_cpu[0].item<float>();\n for (int i = 1; i < num_blocks; i++) {\n float val = block_maxes_cpu[i].item<float>();\n if (val > global_max) global_max = val;\n }\n \n float sum_exp = 0.0f;\n for (int i = 0; i < num_blocks; i++) {\n float max_i = block_maxes_cpu[i].item<float>();\n float sum_i = block_sums_cpu[i].item<float>();\n // Renormalize: exp(max_i) * sum_i -> exp(global_max) * exp(max_i - global_max) * sum_i\n sum_exp += expf(max_i - global_max) * sum_i;\n }\n \n float global_logsumexp = global_max + logf(sum_exp);\n \n // Third pass: normalize and compute weighted mean and ESS\n auto weighted_mean_partial = torch::zeros({num_blocks, dim}, options);\n auto ess_partial = torch::zeros({num_blocks}, options);\n \n normalize_and_reduce_kernel<<<num_blocks, BLOCK_SIZE>>>(\n samples.data_ptr<float>(),\n log_w.data_ptr<float>(),\n global_logsumexp,\n weighted_mean_partial.data_ptr<float>(),\n ess_partial.data_ptr<float>(),\n N, dim\n );\n \n // Final reduction\n auto weighted_mean_out = torch::empty({dim}, options);\n auto ess_out = torch::empty({1}, options);\n \n final_reduce_kernel<<<1, dim>>>(\n weighted_mean_partial.data_ptr<float>(),\n ess_partial.data_ptr<float>(),\n weighted_mean_out.data_ptr<float>(),\n ess_out.data_ptr<float>(),\n num_blocks, dim\n );\n \n return std::make_tuple(weighted_mean_out, ess_out);\n}\n'''\n\ncpp_source = r'''\nstd::tuple<torch::Tensor, torch::Tensor> importance_sampling_cuda(\n torch::Tensor samples,\n torch::Tensor mix_mus,\n torch::Tensor mix_stds,\n torch::Tensor mix_weights,\n torch::Tensor q_mu,\n torch::Tensor q_std);\n'''\n\nmodule = load_inline(\n name='importance_sampling_ext',\n cuda_sources=cuda_source,\n cpp_sources=cpp_source,\n functions=['importance_sampling_cuda'],\n extra_cuda_cflags=['-O3'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n \"\"\"\n Self-normalized importance sampling estimate with CUDA acceleration.\n \"\"\"\n \n def __init__(self, dim, num_components=3):\n super().__init__()\n self.dim = dim\n self.num_components = num_components\n # Target mixture 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 # Proposal\n self.register_buffer(\"q_mu\", torch.zeros(dim))\n self.register_buffer(\"q_std\", torch.ones(dim) * 3.0)\n \n def forward(self, samples: torch.Tensor) -> tuple:\n \"\"\"\n Args:\n samples: (N, dim) \u2013 samples drawn from q\n Returns:\n (weighted_mean, ess):\n weighted_mean: (dim,) \u2013 importance-weighted mean estimate of E_p[x]\n ess: scalar \u2013 effective sample size\n \"\"\"\n # Ensure contiguous and on same device\n samples = samples.contiguous()\n mix_mus = self.mix_mus.contiguous()\n mix_stds = self.mix_stds.contiguous()\n mix_weights = self.mix_weights.contiguous()\n q_mu = self.q_mu.contiguous()\n q_std = self.q_std.contiguous()\n \n weighted_mean, ess = module.importance_sampling_cuda(\n samples, mix_mus, mix_stds, mix_weights, q_mu, q_std\n )\n \n return weighted_mean, ess.squeeze()"
}