call → compile_kernel
callcompile_kernel
{
"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#include <cmath>\n\n// Constants for log calculations\nconstexpr float LOG_2PI = 1.8378770664093453f; // log(2*pi)\nconstexpr float HALF_LOG_2PI = 0.9189385332046727f; // 0.5 * log(2*pi)\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// Compute softmax of logits\n__global__ void softmax_kernel(const float* logits, float* weights, int K) {\n extern __shared__ float smem[];\n \n // Find max for numerical stability\n float max_val = -FLT_MAX;\n for (int i = threadIdx.x; i < K; i += blockDim.x) {\n max_val = fmaxf(max_val, logits[i]);\n }\n \n // Block reduce for max\n smem[threadIdx.x] = max_val;\n __syncthreads();\n for (int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (threadIdx.x < s) {\n smem[threadIdx.x] = fmaxf(smem[threadIdx.x], smem[threadIdx.x + s]);\n }\n __syncthreads();\n }\n max_val = smem[0];\n \n // Compute exp and sum\n float sum = 0.0f;\n for (int i = threadIdx.x; i < K; i += blockDim.x) {\n sum += expf(logits[i] - max_val);\n }\n \n // Block reduce for sum\n smem[threadIdx.x] = sum;\n __syncthreads();\n for (int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (threadIdx.x < s) {\n smem[threadIdx.x] += smem[threadIdx.x + s];\n }\n __syncthreads();\n }\n sum = smem[0];\n \n // Write normalized weights\n for (int i = threadIdx.x; i < K; i += blockDim.x) {\n weights[i] = expf(logits[i] - max_val) / sum;\n }\n}\n\n// Sample from categorical distribution using inverse CDF\n__device__ int sample_categorical(float* probs, int K, float u) {\n float cumsum = 0.0f;\n for (int i = 0; i < K; i++) {\n cumsum += probs[i];\n if (u < cumsum) {\n return i;\n }\n }\n return K - 1; // Fallback\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__ weights, // (K,)\n float* __restrict__ z, // (B, D)\n float* __restrict__ log_q_z, // (B,)\n int* __restrict__ components, // (B,) - sampled components\n curandState* states, // (B,)\n int B, int K, int D\n) {\n int b = blockIdx.x; // One block per batch element\n \n if (b >= B) return;\n \n // Each thread handles one batch element\n // Load curand state\n curandState local_state = states[b];\n \n // Sample component using inverse CDF\n float u = curand_uniform(&local_state);\n int comp = 0;\n float cumsum = 0.0f;\n for (int k = 0; k < K; k++) {\n cumsum += weights[k];\n if (u < cumsum) {\n comp = k;\n break;\n }\n }\n components[b] = comp;\n \n // Sample z for each dimension\n for (int d = threadIdx.x; d < D; d += blockDim.x) {\n int idx = b * K * D + comp * D + d;\n float mu_val = mu[idx];\n float log_var_val = log_var[idx];\n float std_val = expf(0.5f * log_var_val);\n \n // Box-Muller transform for normal random\n float u1 = curand_uniform(&local_state);\n float u2 = curand_uniform(&local_state);\n // Avoid log(0)\n while (u1 <= 0.0f) u1 = curand_uniform(&local_state);\n float rand_normal = sqrtf(-2.0f * logf(u1)) * cosf(2.0f * 3.14159265f * u2);\n \n z[b * D + d] = mu_val + std_val * rand_normal;\n }\n \n // Save state back\n if (threadIdx.x == 0) {\n states[b] = local_state;\n }\n __syncthreads();\n \n // Compute log q(z) under mixture\n // log_probs = -0.5 * ((z - mu) / std)^2 - 0.5 * log_var - 0.5 * log(2*pi)\n // Use shared memory for per-component sums\n extern __shared__ float smem[];\n float* comp_sums = smem; // K elements\n float* thread_sums = &smem[K]; // blockDim.x elements\n \n // Each thread computes partial sums for all components\n for (int k = threadIdx.x; k < K; k += blockDim.x) {\n comp_sums[k] = 0.0f;\n }\n __syncthreads();\n \n // Each thread computes contribution for dimensions it handles\n float my_sum[10]; // Max K components, assume K <= 10\n for (int k = 0; k < K; k++) {\n my_sum[k] = 0.0f;\n }\n \n for (int d = threadIdx.x; d < D; d += blockDim.x) {\n float z_val = z[b * D + d];\n for (int k = 0; k < K; k++) {\n int idx = b * K * D + k * D + d;\n float mu_val = mu[idx];\n float log_var_val = log_var[idx];\n float std_val = expf(0.5f * log_var_val);\n float diff = (z_val - mu_val) / std_val;\n float log_prob_d = -0.5f * diff * diff - 0.5f * log_var_val - HALF_LOG_2PI;\n my_sum[k] += log_prob_d;\n }\n }\n \n // Reduce within block for each component\n for (int k = 0; k < K; k++) {\n thread_sums[threadIdx.x] = my_sum[k];\n __syncthreads();\n for (int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (threadIdx.x < s) {\n thread_sums[threadIdx.x] += thread_sums[threadIdx.x + s];\n }\n __syncthreads();\n }\n if (threadIdx.x == 0) {\n comp_sums[k] = thread_sums[0] + logf(weights[k]);\n }\n __syncthreads();\n }\n \n // Compute logsumexp using parallel reduction\n if (threadIdx.x == 0) {\n float max_val = -FLT_MAX;\n for (int k = 0; k < K; k++) {\n max_val = fmaxf(max_val, comp_sums[k]);\n }\n \n float sum_exp = 0.0f;\n for (int k = 0; k < K; k++) {\n sum_exp += expf(comp_sums[k] - max_val);\n }\n log_q_z[b] = max_val + logf(sum_exp);\n }\n}\n\n// Simple reparameterization kernel (no mixture)\n__global__ void simple_reparam_kernel(\n const float* __restrict__ mu, // (B, D)\n const float* __restrict__ log_var, // (B, D)\n float* __restrict__ z, // (B, D)\n float* __restrict__ log_q_z, // (B,)\n curandState* states, // (B,)\n int B, int D\n) {\n int b = blockIdx.x; // One block per batch element\n \n if (b >= B) return;\n \n curandState local_state = states[b];\n \n // Shared memory for reduction\n extern __shared__ float smem[];\n float& thread_sum = smem[threadIdx.x];\n \n thread_sum = 0.0f;\n \n for (int d = threadIdx.x; d < D; d += blockDim.x) {\n int idx = b * D + d;\n float mu_val = mu[idx];\n float log_var_val = log_var[idx];\n float std_val = expf(0.5f * log_var_val);\n \n // Box-Muller transform\n float u1 = curand_uniform(&local_state);\n float u2 = curand_uniform(&local_state);\n while (u1 <= 0.0f) u1 = curand_uniform(&local_state);\n float rand_normal = sqrtf(-2.0f * logf(u1)) * cosf(2.0f * 3.14159265f * u2);\n \n float z_val = mu_val + std_val * rand_normal;\n z[idx] = z_val;\n \n // Compute log prob contribution\n float diff = (z_val - mu_val) / std_val;\n float log_prob_d = -0.5f * diff * diff - 0.5f * log_var_val - HALF_LOG_2PI;\n thread_sum += log_prob_d;\n }\n \n // Save state\n if (threadIdx.x == 0) {\n states[b] = local_state;\n }\n \n __syncthreads();\n \n // Block reduce\n for (int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (threadIdx.x < s) {\n smem[threadIdx.x] += smem[threadIdx.x + s];\n }\n __syncthreads();\n }\n \n if (threadIdx.x == 0) {\n log_q_z[b] = smem[0];\n }\n}\n\n// Host functions\nvoid init_curand_cuda(curandState* states, unsigned long long seed, int n) {\n int threads = 256;\n int blocks = (n + threads - 1) / threads;\n init_curand_kernel<<<blocks, threads>>>(states, seed, n);\n}\n\nvoid softmax_cuda(const float* logits, float* weights, int K) {\n int threads = 256;\n softmax_kernel<<<1, threads, threads * sizeof(float)>>>(logits, weights, K);\n}\n\nvoid mixture_reparam_cuda(\n const float* mu, const float* log_var, const float* weights,\n float* z, float* log_q_z, int* components,\n curandState* states, int B, int K, int D\n) {\n int threads = 256;\n size_t shared_mem = (K + threads) * sizeof(float);\n mixture_reparam_kernel<<<B, threads, shared_mem>>>(\n mu, log_var, weights, z, log_q_z, components, states, B, K, D\n );\n}\n\nvoid simple_reparam_cuda(\n const float* mu, const float* log_var,\n float* z, float* log_q_z,\n curandState* states, int B, int D\n) {\n int threads = 256;\n size_t shared_mem = threads * sizeof(float);\n simple_reparam_kernel<<<B, threads, shared_mem>>>(\n mu, log_var, z, log_q_z, states, B, D\n );\n}\n\"\"\"\n\ncpp_source = \"\"\"\n#include <torch/extension.h>\n#include <curand_kernel.h>\n\n// Declare functions\nvoid init_curand_cuda(curandState* states, unsigned long long seed, int n);\nvoid softmax_cuda(const float* logits, float* weights, int K);\nvoid mixture_reparam_cuda(\n const float* mu, const float* log_var, const float* weights,\n float* z, float* log_q_z, int* components,\n curandState* states, int B, int K, int D\n);\nvoid simple_reparam_cuda(\n const float* mu, const float* log_var,\n float* z, float* log_q_z,\n curandState* states, int B, int D\n);\n\n// Wrapper functions for PyTorch\nvoid init_curand(torch::Tensor states, unsigned long long seed, int n) {\n init_curand_cuda(reinterpret_cast<curandState*>(states.data_ptr()), seed, n);\n}\n\nvoid softmax_weights(torch::Tensor logits, torch::Tensor weights, int K) {\n softmax_cuda(logits.data_ptr<float>(), weights.data_ptr<float>(), K);\n}\n\nvoid mixture_reparam(\n torch::Tensor mu, torch::Tensor log_var, torch::Tensor weights,\n torch::Tensor z, torch::Tensor log_q_z, torch::Tensor components,\n torch::Tensor states, int B, int K, int D\n) {\n mixture_reparam_cuda(\n mu.data_ptr<float>(), log_var.data_ptr<float>(), weights.data_ptr<float>(),\n z.data_ptr<float>(), log_q_z.data_ptr<float>(), components.data_ptr<int>(),\n reinterpret_cast<curandState*>(states.data_ptr()), B, K, D\n );\n}\n\nvoid simple_reparam(\n torch::Tensor mu, torch::Tensor log_var,\n torch::Tensor z, torch::Tensor log_q_z,\n torch::Tensor states, int B, int D\n) {\n simple_reparam_cuda(\n mu.data_ptr<float>(), log_var.data_ptr<float>(),\n z.data_ptr<float>(), log_q_z.data_ptr<float>(),\n reinterpret_cast<curandState*>(states.data_ptr()), B, D\n );\n}\n\"\"\"\n\nmodule = load_inline(\n name='reparam_cuda',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['init_curand', 'softmax_weights', 'mixture_reparam', 'simple_reparam'],\n extra_cuda_cflags=['-O3', '--use_fast_math', '-I/usr/local/cuda/include'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n \"\"\"\n Reparameterization trick for differentiable sampling from a diagonal\n Gaussian. Uses custom CUDA kernels for better performance.\n \"\"\"\n\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 # Allocate cuRAND states (one per batch element, max batch 1024)\n self.max_batch = 1024\n self.curand_states = torch.empty(self.max_batch * 256, dtype=torch.uint8, device='cuda')\n module.init_curand(self.curand_states, 42, self.max_batch)\n self._seed = 42\n\n def forward(self, mu: torch.Tensor, log_var: torch.Tensor) -> tuple:\n B = mu.shape[0]\n K = self.num_components\n D = self.latent_dim\n \n # Ensure inputs are on GPU\n if not mu.is_cuda:\n mu = mu.cuda()\n log_var = log_var.cuda()\n \n # Ensure contiguous\n mu = mu.contiguous()\n log_var = log_var.contiguous()\n \n # Re-initialize curand states if batch size changed\n if B > self.max_batch:\n self.max_batch = B * 2\n self.curand_states = torch.empty(self.max_batch * 256, dtype=torch.uint8, device='cuda')\n self._seed += 1\n module.init_curand(self.curand_states, self._seed, self.max_batch)\n \n if K > 1 and mu.dim() == 3:\n # Mixture of Gaussians case\n weights = torch.softmax(self.logits, dim=0)\n \n z = torch.empty(B, D, dtype=mu.dtype, device=mu.device)\n log_q_z = torch.empty(B, dtype=mu.dtype, device=mu.device)\n components = torch.empty(B, dtype=torch.int32, device=mu.device)\n \n module.mixture_reparam(\n mu.float(), log_var.float(), weights.float(),\n z.float(), log_q_z.float(), components.int(),\n self.curand_states, B, K, D\n )\n else:\n # Simple case\n if mu.dim() == 3:\n mu = mu[:, 0]\n log_var = log_var[:, 0]\n \n z = torch.empty(B, D, dtype=mu.dtype, device=mu.device)\n log_q_z = torch.empty(B, dtype=mu.dtype, device=mu.device)\n \n module.simple_reparam(\n mu.float(), log_var.float(),\n z.float(), log_q_z.float(),\n self.curand_states, B, D\n )\n \n return z, log_q_z"
}