call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\nimport math\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <curand_kernel.h>\n#include <cfloat>\n\n// Constants\nconstexpr float PI = 3.14159265358979323846f;\nconstexpr float NEG_HALF_LOG_2PI = -0.9189385332046727f; // -0.5 * log(2*pi)\n\n// Warp-level reduction for sum\n__inline__ __device__ float warp_reduce_sum(float val) {\n for (int offset = 16; offset > 0; offset /= 2) {\n val += __shfl_down_sync(0xffffffff, val, offset);\n }\n return val;\n}\n\n// Warp-level reduce max\n__inline__ __device__ float warp_reduce_max(float val) {\n for (int offset = 16; offset > 0; offset /= 2) {\n val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));\n }\n return val;\n}\n\n// Block-level reduce max\n__device__ float block_reduce_max(float val) {\n __shared__ float shared[32];\n int lane = threadIdx.x % 32;\n int wid = threadIdx.x / 32;\n \n val = warp_reduce_max(val);\n if (lane == 0) {\n shared[wid] = val;\n }\n __syncthreads();\n \n if (threadIdx.x == 0) {\n val = shared[0];\n for (int i = 1; i < blockDim.x / 32; ++i) {\n val = fmaxf(val, shared[i]);\n }\n shared[0] = val;\n }\n __syncthreads();\n return shared[0];\n}\n\n// Block-level reduce sum\n__device__ float block_reduce_sum(float val) {\n __shared__ float shared[32];\n int lane = threadIdx.x % 32;\n int wid = threadIdx.x / 32;\n \n val = warp_reduce_sum(val);\n if (lane == 0) {\n shared[wid] = val;\n }\n __syncthreads();\n \n if (threadIdx.x == 0) {\n val = shared[0];\n for (int i = 1; i < blockDim.x / 32; ++i) {\n val += shared[i];\n }\n shared[0] = val;\n }\n __syncthreads();\n return shared[0];\n}\n\n// Initialize cuRAND states\n__global__ void init_curand_kernel(curandState* states, unsigned long long seed, int n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < n) {\n curand_init(seed, idx, 0, &states[idx]);\n }\n}\n\n// Main kernel for mixture of Gaussians reparameterization\n__global__ void mixture_reparam_kernel(\n const float* __restrict__ mu, // (B, K, D)\n const float* __restrict__ log_var, // (B, K, D)\n const float* __restrict__ logits, // (K,)\n float* __restrict__ z, // (B, D)\n float* __restrict__ log_q_z, // (B,)\n curandState* states, // (B,)\n int batch_size,\n int num_components,\n int latent_dim\n) {\n int b = blockIdx.x;\n if (b >= batch_size) return;\n \n int tid = threadIdx.x;\n \n // Load cuRAND state\n curandState local_state = states[b];\n \n // Step 1: Compute softmax weights in shared memory\n extern __shared__ float smem[];\n float* weights = smem; // size K\n float* log_weights = smem + num_components; // size K\n float* log_probs_comp = smem + 2 * num_components; // size K (for logsumexp reduction)\n \n // Compute max logit for numerical stability\n float max_logit = -FLT_MAX;\n for (int k = tid; k < num_components; k += blockDim.x) {\n max_logit = fmaxf(max_logit, logits[k]);\n }\n max_logit = block_reduce_max(max_logit);\n \n // Compute softmax weights\n float sum_exp = 0.0f;\n for (int k = tid; k < num_components; k += blockDim.x) {\n float exp_val = expf(logits[k] - max_logit);\n weights[k] = exp_val;\n sum_exp += exp_val;\n }\n sum_exp = block_reduce_sum(sum_exp);\n \n for (int k = tid; k < num_components; k += blockDim.x) {\n weights[k] /= sum_exp;\n log_weights[k] = logf(weights[k] + 1e-20f); // avoid log(0)\n }\n __syncthreads();\n \n // Step 2: Sample component using multinomial (cumulative sum method)\n int sampled_comp = 0;\n if (tid == 0) {\n float u = curand_uniform(&local_state);\n float cumsum = 0.0f;\n for (int k = 0; k < num_components; ++k) {\n cumsum += weights[k];\n if (u <= cumsum) {\n sampled_comp = k;\n break;\n }\n }\n }\n __syncthreads();\n \n // Broadcast sampled component to all threads\n __shared__ int shared_comp;\n if (tid == 0) {\n shared_comp = sampled_comp;\n }\n __syncthreads();\n sampled_comp = shared_comp;\n \n // Step 3: Compute z = mu_k + std_k * eps for sampled component k\n // Each thread handles multiple latent dimensions\n float* z_out = z + b * latent_dim;\n const float* mu_b = mu + b * num_components * latent_dim;\n const float* log_var_b = log_var + b * num_components * latent_dim;\n \n for (int d = tid; d < latent_dim; d += blockDim.x) {\n float mu_val = mu_b[sampled_comp * latent_dim + d];\n float log_var_val = log_var_b[sampled_comp * latent_dim + d];\n float std_val = expf(0.5f * log_var_val);\n float eps = curand_normal(&local_state);\n z_out[d] = mu_val + std_val * eps;\n }\n __syncthreads();\n \n // Step 4: Compute log q(z) under mixture\n // log q(z) = logsumexp_k [ log w_k + sum_d log N(z_d | mu_{k,d}, var_{k,d}) ]\n // log N(z|mu,sigma^2) = -0.5 * ((z-mu)/sigma)^2 - 0.5 * log(var) - 0.5 * log(2*pi)\n \n // For each component, compute sum over latent dims\n // Each block thread works on a subset of latent dims per component\n \n __shared__ float log_probs_by_comp[32]; // max 32 components for simplicity\n \n // Initialize component log probs\n for (int k = tid; k < num_components; k += blockDim.x) {\n log_probs_by_comp[k] = log_weights[k]; // Start with log(w_k)\n }\n __syncthreads();\n \n // Accumulate per-component log probabilities\n for (int k = 0; k < num_components; ++k) {\n float local_sum = 0.0f;\n const float* mu_k = mu_b + k * latent_dim;\n const float* log_var_k = log_var_b + k * latent_dim;\n \n for (int d = tid; d < latent_dim; d += blockDim.x) {\n float z_val = z_out[d];\n float mu_val = mu_k[d];\n float log_var_val = log_var_k[d];\n float std_val = expf(0.5f * log_var_val);\n \n // Log prob contribution: -0.5*((z-mu)/std)^2 - 0.5*log_var - 0.5*log(2pi)\n float diff = (z_val - mu_val) / std_val;\n float log_prob_d = -0.5f * diff * diff - 0.5f * log_var_val + NEG_HALF_LOG_2PI;\n local_sum += log_prob_d;\n }\n \n local_sum = block_reduce_sum(local_sum);\n if (tid == 0) {\n log_probs_by_comp[k] += local_sum;\n }\n __syncthreads();\n }\n \n // Compute logsumexp over components\n if (tid == 0) {\n float max_log_prob = log_probs_by_comp[0];\n for (int k = 1; k < num_components; ++k) {\n max_log_prob = fmaxf(max_log_prob, log_probs_by_comp[k]);\n }\n \n float sum_exp_log = 0.0f;\n for (int k = 0; k < num_components; ++k) {\n sum_exp_log += expf(log_probs_by_comp[k] - max_log_prob);\n }\n \n log_q_z[b] = max_log_prob + logf(sum_exp_log);\n }\n \n // Save state back\n if (tid == 0) {\n states[b] = local_state;\n }\n}\n\n// Single component case kernel\n__global__ void single_component_kernel(\n const float* __restrict__ mu, // (B, D) or (B, 1, D)\n const float* __restrict__ log_var, // (B, D) or (B, 1, D)\n float* __restrict__ z, // (B, D)\n float* __restrict__ log_q_z, // (B,)\n curandState* states, // (B,)\n int batch_size,\n int latent_dim,\n bool has_components_dim\n) {\n int b = blockIdx.x;\n if (b >= batch_size) return;\n \n int tid = threadIdx.x;\n curandState local_state = states[b];\n \n float* z_out = z + b * latent_dim;\n const float* mu_b = has_components_dim ? mu + b * latent_dim : mu + b * latent_dim;\n const float* log_var_b = has_components_dim ? log_var + b * latent_dim : log_var + b * latent_dim;\n \n float local_log_q = 0.0f;\n \n for (int d = tid; d < latent_dim; d += blockDim.x) {\n float mu_val = mu_b[d];\n float log_var_val = log_var_b[d];\n float std_val = expf(0.5f * log_var_val);\n float eps = curand_normal(&local_state);\n float z_val = mu_val + std_val * eps;\n z_out[d] = z_val;\n \n // Accumulate log prob\n float diff = (z_val - mu_val) / std_val;\n float log_prob_d = -0.5f * diff * diff - 0.5f * log_var_val + NEG_HALF_LOG_2PI;\n local_log_q += log_prob_d;\n }\n \n local_log_q = block_reduce_sum(local_log_q);\n \n if (tid == 0) {\n log_q_z[b] = local_log_q;\n states[b] = local_state;\n }\n}\n\n// Host functions\nvoid init_curand_cuda(torch::Tensor states, unsigned long long seed, int n) {\n curandState* states_ptr = reinterpret_cast<curandState*>(states.data_ptr());\n int threads = 256;\n int blocks = (n + threads - 1) / threads;\n init_curand_kernel<<<blocks, threads>>>(states_ptr, seed, n);\n}\n\nvoid mixture_reparam_cuda(\n torch::Tensor mu,\n torch::Tensor log_var,\n torch::Tensor logits,\n torch::Tensor z,\n torch::Tensor log_q_z,\n torch::Tensor states,\n int batch_size,\n int num_components,\n int latent_dim\n) {\n const float* mu_ptr = mu.data_ptr<float>();\n const float* log_var_ptr = log_var.data_ptr<float>();\n const float* logits_ptr = logits.data_ptr<float>();\n float* z_ptr = z.data_ptr<float>();\n float* log_q_z_ptr = log_q_z.data_ptr<float>();\n curandState* states_ptr = reinterpret_cast<curandState*>(states.data_ptr());\n \n int threads = 256;\n size_t shared_mem = 3 * num_components * sizeof(float);\n \n mixture_reparam_kernel<<<batch_size, threads, shared_mem>>>(\n mu_ptr, log_var_ptr, logits_ptr, z_ptr, log_q_z_ptr, states_ptr,\n batch_size, num_components, latent_dim\n );\n}\n\nvoid single_component_cuda(\n torch::Tensor mu,\n torch::Tensor log_var,\n torch::Tensor z,\n torch::Tensor log_q_z,\n torch::Tensor states,\n int batch_size,\n int latent_dim,\n bool has_components_dim\n) {\n const float* mu_ptr = mu.data_ptr<float>();\n const float* log_var_ptr = log_var.data_ptr<float>();\n float* z_ptr = z.data_ptr<float>();\n float* log_q_z_ptr = log_q_z.data_ptr<float>();\n curandState* states_ptr = reinterpret_cast<curandState*>(states.data_ptr());\n \n int threads = 256;\n single_component_kernel<<<batch_size, threads>>>(\n mu_ptr, log_var_ptr, z_ptr, log_q_z_ptr, states_ptr,\n batch_size, latent_dim, has_components_dim\n );\n}\n\"\"\"\n\ncpp_source = \"\"\"\n#include <torch/extension.h>\n#include <curand_kernel.h>\n\n// Size of curandState in bytes\nconstexpr size_t CURAND_STATE_SIZE = sizeof(curandState);\n\nvoid init_curand_cuda(torch::Tensor states, unsigned long long seed, int n);\nvoid mixture_reparam_cuda(torch::Tensor mu, torch::Tensor log_var, torch::Tensor logits, \n torch::Tensor z, torch::Tensor log_q_z, torch::Tensor states,\n int batch_size, int num_components, int latent_dim);\nvoid single_component_cuda(torch::Tensor mu, torch::Tensor log_var,\n torch::Tensor z, torch::Tensor log_q_z, torch::Tensor states,\n int batch_size, int latent_dim, bool has_components_dim);\n\ntorch::Tensor create_curand_states(int batch_size, unsigned long long seed) {\n // Allocate raw bytes for curand states\n auto options = torch::TensorOptions().dtype(torch::kByte).device(torch::kCUDA);\n torch::Tensor states = torch::empty({static_cast<long>(batch_size * CURAND_STATE_SIZE)}, options);\n init_curand_cuda(states, seed, batch_size);\n return states;\n}\n\nstd::tuple<torch::Tensor, torch::Tensor> forward_mixture(\n torch::Tensor mu, torch::Tensor log_var, torch::Tensor logits, torch::Tensor states\n) {\n int batch_size = mu.size(0);\n int num_components = mu.size(1);\n int latent_dim = mu.size(2);\n \n auto z = torch::empty({batch_size, latent_dim}, mu.options());\n auto log_q_z = torch::empty({batch_size}, mu.options());\n \n mixture_reparam_cuda(mu, log_var, logits, z, log_q_z, states,\n batch_size, num_components, latent_dim);\n \n return std::make_tuple(z, log_q_z);\n}\n\nstd::tuple<torch::Tensor, torch::Tensor> forward_single(\n torch::Tensor mu, torch::Tensor log_var, torch::Tensor states, bool has_components_dim\n) {\n int batch_size = mu.size(0);\n int latent_dim = has_components_dim ? mu.size(2) : mu.size(1);\n \n auto z = torch::empty({batch_size, latent_dim}, mu.options());\n auto z = torch::empty({batch_size, latent_dim}, mu.options());\n auto log_q_z = torch::empty({batch_size}, mu.options());\n \n single_component_cuda(mu, log_var, z, log_q_z, states,\n batch_size, latent_dim, has_components_dim);\n \n return std::make_tuple(z, log_q_z);\n}\n\"\"\"\n\n# Load the extension\nextension = load_inline(\n name='reparam_cuda',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['create_curand_states', 'forward_mixture', 'forward_single'],\n extra_cuda_cflags=['-O3', '--use_fast_math', '-std=c++17'],\n verbose=False\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self, latent_dim, num_components=1):\n super().__init__()\n self.latent_dim = latent_dim\n self.num_components = num_components\n if num_components > 1:\n self.logits = nn.Parameter(torch.zeros(num_components))\n \n # Initialize curand states on module creation\n self.register_buffer('_curand_states', None, persistent=False)\n self._batch_size_cached = 0\n \n def _ensure_curand_states(self, batch_size, device):\n if self._curand_states is None or self._batch_size_cached != batch_size:\n import random\n seed = random.randint(0, 2**31 - 1)\n self._curand_states = extension.create_curand_states(batch_size, seed)\n self._batch_size_cached = batch_size\n \n def forward(self, mu: torch.Tensor, log_var: torch.Tensor) -> tuple:\n batch_size = mu.shape[0]\n device = mu.device\n \n # Ensure contiguous float32 tensors\n mu = mu.contiguous().float()\n log_var = log_var.contiguous().float()\n \n self._ensure_curand_states(batch_size, device)\n \n if self.num_components > 1 and mu.dim() == 3:\n z, log_q_z = extension.forward_mixture(\n mu, log_var, self.logits.float().contiguous(), self._curand_states\n )\n else:\n # Single component case\n has_components_dim = (mu.dim() == 3)\n if has_components_dim:\n mu = mu[:, 0].contiguous()\n log_var = log_var[:, 0].contiguous()\n z, log_q_z = extension.forward_single(\n mu, log_var, self._curand_states, has_components_dim\n )\n \n return z, log_q_z"
}