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 = \"\"\"\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n// Constants\nconstexpr int BLOCK_SIZE = 256;\nconstexpr float NEG_INF = -1e30f;\n\n// Warp-level logsumexp reduction\n__device__ float warp_logsumexp(float val, int width) {\n for (int offset = width / 2; offset > 0; offset /= 2) {\n float other = __shfl_down_sync(0xffffffff, val, offset);\n float max_val = fmaxf(val, other);\n val = max_val + logf(expf(val - max_val) + expf(other - max_val));\n }\n return val;\n}\n\n// Block-level logsumexp using shared memory\n__device__ float block_logsumexp(float val, float* shmem) {\n int lane = threadIdx.x % 32;\n int warp_id = threadIdx.x / 32;\n \n // Warp reduction\n val = warp_logsumexp(val, 32);\n \n // First lane of each warp writes to shared memory\n if (lane == 0) {\n shmem[warp_id] = val;\n }\n __syncthreads();\n \n // First warp does final reduction\n if (warp_id == 0) {\n val = (lane < (BLOCK_SIZE / 32)) ? shmem[lane] : NEG_INF;\n val = warp_logsumexp(val, BLOCK_SIZE / 32);\n if (lane == 0) {\n shmem[0] = val;\n }\n }\n __syncthreads();\n \n return shmem[0];\n}\n\n// Normalize a 1D log vector (log_init: S -> S normalized)\n__global__ void normalize_log_1d_kernel(const float* __restrict__ input, float* __restrict__ output, int size) {\n float local_max = NEG_INF;\n \n // Find max using block reduction\n for (int i = threadIdx.x; i < size; i += blockDim.x) {\n local_max = fmaxf(local_max, input[i]);\n }\n \n __shared__ float shmem[BLOCK_SIZE / 32];\n float max_val = block_logsumexp(local_max, shmem);\n // Actually need just max for logsumexp, but let's use the logsumexp pattern\n \n // Actually simpler: compute logsumexp directly\n // Reset and compute logsumexp\n local_max = NEG_INF;\n for (int i = threadIdx.x; i < size; i += blockDim.x) {\n local_max = fmaxf(local_max, input[i]);\n }\n max_val = block_logsumexp(local_max, shmem);\n \n // Now compute logsumexp with correct max\n float local_sum = NEG_INF;\n for (int i = threadIdx.x; i < size; i += blockDim.x) {\n float contrib = input[i] - max_val;\n if (local_sum == NEG_INF) {\n local_sum = contrib;\n } else {\n local_sum = local_sum + logf(expf(local_sum) + expf(contrib));\n }\n }\n \n float log_sum = block_logsumexp(local_sum, shmem);\n float logsumexp_val = max_val + log_sum;\n \n // Normalize\n for (int i = threadIdx.x; i < size; i += blockDim.x) {\n output[i] = input[i] - logsumexp_val;\n }\n}\n\n// Normalize a 2D log matrix along dim=-1 (rows: SxS -> SxS normalized)\n__global__ void normalize_log_2d_kernel(const float* __restrict__ input, float* __restrict__ output, int rows, int cols) {\n int row = blockIdx.x;\n if (row >= rows) return;\n \n const float* row_in = input + row * cols;\n float* row_out = output + row * cols;\n \n // Find max\n float local_max = NEG_INF;\n for (int j = threadIdx.x; j < cols; j += blockDim.x) {\n local_max = fmaxf(local_max, row_in[j]);\n }\n \n __shared__ float shmem[BLOCK_SIZE / 32];\n float max_val = block_logsumexp(local_max, shmem);\n \n // Compute logsumexp\n float local_sum = NEG_INF;\n for (int j = threadIdx.x; j < cols; j += blockDim.x) {\n float contrib = row_in[j] - max_val;\n if (local_sum == NEG_INF) {\n local_sum = contrib;\n } else {\n local_sum = local_sum + logf(expf(local_sum) + expf(contrib));\n }\n }\n \n float log_sum = block_logsumexp(local_sum, shmem);\n float logsumexp_val = max_val + log_sum;\n \n // Normalize\n for (int j = threadIdx.x; j < cols; j += blockDim.x) {\n row_out[j] = row_in[j] - logsumexp_val;\n }\n}\n\n// Main forward algorithm kernel\n// observations: (B, T), log_pi: (S,), log_A: (S, S), log_B: (S, num_obs), log_alpha_out: (B, S)\n__global__ void hmm_forward_kernel(\n const int* __restrict__ observations,\n const float* __restrict__ log_pi,\n const float* __restrict__ log_A,\n const float* __restrict__ log_B,\n float* __restrict__ log_alpha,\n float* __restrict__ log_lik_out,\n int B, int T, int S, int num_obs\n) {\n int b = blockIdx.x;\n if (b >= B) return;\n \n const int* obs_b = observations + b * T;\n float* alpha_b = log_alpha + b * S;\n \n // Initialize: log_alpha[j] = log_pi[j] + log_B[j, obs[0]]\n int obs_0 = obs_b[0];\n for (int j = threadIdx.x; j < S; j += blockDim.x) {\n alpha_b[j] = log_pi[j] + log_B[j * num_obs + obs_0];\n }\n __syncthreads();\n \n // Iterate through time steps\n for (int t = 1; t < T; t++) {\n int obs_t = obs_b[t];\n \n // For each j: log_alpha_new[j] = logsumexp_i(log_alpha[i] + log_A[i,j]) + log_B[j, obs_t]\n for (int j = threadIdx.x; j < S; j += blockDim.x) {\n float max_val = NEG_INF;\n // Find max\n for (int i = 0; i < S; i++) {\n float val = alpha_b[i] + log_A[i * S + j];\n max_val = fmaxf(max_val, val);\n }\n \n // Compute logsumexp\n float sum_exp = 0.0f;\n for (int i = 0; i < S; i++) {\n float val = alpha_b[i] + log_A[i * S + j];\n sum_exp += expf(val - max_val);\n }\n float log_alpha_new = max_val + logf(sum_exp) + log_B[j * num_obs + obs_t];\n \n // Store in alpha_b for next iteration (we need double buffer or careful sync)\n // Use shared memory as intermediate buffer\n extern __shared__ float temp_alpha[];\n temp_alpha[j] = log_alpha_new;\n }\n __syncthreads();\n \n // Copy back\n for (int j = threadIdx.x; j < S; j += blockDim.x) {\n alpha_b[j] = temp_alpha[j];\n }\n __syncthreads();\n }\n \n // Final logsumexp to get log_lik\n float local_max = NEG_INF;\n for (int j = threadIdx.x; j < S; j += blockDim.x) {\n local_max = fmaxf(local_max, alpha_b[j]);\n }\n \n __shared__ float shmem[BLOCK_SIZE / 32 + 1];\n float max_val = block_logsumexp(local_max, shmem);\n \n // This gives wrong result, need proper max reduction\n // Let's do it simpler\n __shared__ float shared_max;\n if (threadIdx.x == 0) shared_max = NEG_INF;\n __syncthreads();\n \n for (int j = threadIdx.x; j < S; j += blockDim.x) {\n atomicMax(reinterpret_cast<int*>(&shared_max), __float_as_int(alpha_b[j]));\n }\n __syncthreads();\n \n max_val = shared_max;\n \n float local_sum = 0.0f;\n for (int j = threadIdx.x; j < S; j += blockDim.x) {\n local_sum += expf(alpha_b[j] - max_val);\n }\n \n __shared__ float total_sum;\n if (threadIdx.x == 0) total_sum = 0.0f;\n __syncthreads();\n \n // Add using atomic add\n atomicAdd(&total_sum, local_sum);\n __syncthreads();\n \n if (threadIdx.x == 0) {\n log_lik_out[b] = max_val + logf(total_sum);\n }\n}\n\ntorch::Tensor hmm_forward_cuda(\n torch::Tensor observations,\n torch::Tensor log_init,\n torch::Tensor log_trans,\n torch::Tensor log_emit\n) {\n int B = observations.size(0);\n int T = observations.size(1);\n int S = log_init.size(0);\n int num_obs = log_emit.size(1);\n \n // Allocate output\n auto options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);\n auto log_alpha = torch::empty({B, S}, options);\n auto log_lik = torch::empty({B}, options);\n \n // Normalize parameters\n auto log_pi = torch::empty({S}, options);\n auto log_A = torch::empty({S, S}, options);\n auto log_B = torch::empty({S, num_obs}, options);\n \n // Launch normalization kernels\n normalize_log_1d_kernel<<<1, BLOCK_SIZE>>>(\n log_init.data_ptr<float>(), log_pi.data_ptr<float>(), S);\n \n normalize_log_2d_kernel<<<S, BLOCK_SIZE>>>(\n log_trans.data_ptr<float>(), log_A.data_ptr<float>(), S, S);\n \n normalize_log_2d_kernel<<<S, BLOCK_SIZE>>>(\n log_emit.data_ptr<float>(), log_B.data_ptr<float>(), S, num_obs);\n \n // Main forward kernel\n int shared_mem_size = S * sizeof(float);\n hmm_forward_kernel<<<B, BLOCK_SIZE, shared_mem_size>>>(\n observations.data_ptr<int>(),\n log_pi.data_ptr<float>(),\n log_A.data_ptr<float>(),\n log_B.data_ptr<float>(),\n log_alpha.data_ptr<float>(),\n log_lik.data_ptr<float>(),\n B, T, S, num_obs\n );\n \n return log_lik;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor hmm_forward_cuda(torch::Tensor observations, torch::Tensor log_init, torch::Tensor log_trans, torch::Tensor log_emit);\n\"\"\"\n\nhmm_module = load_inline(\n name='hmm_forward',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['hmm_forward_cuda'],\n extra_cuda_cflags=['-O3'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self, num_states, num_obs):\n super().__init__()\n self.num_states = num_states\n self.num_obs = num_obs\n \n self.log_init = nn.Parameter(torch.zeros(num_states))\n self.log_trans = nn.Parameter(torch.zeros(num_states, num_states))\n self.log_emit = nn.Parameter(torch.zeros(num_states, num_obs))\n \n def forward(self, observations: torch.Tensor) -> torch.Tensor:\n return hmm_module.hmm_forward_cuda(\n observations,\n self.log_init,\n self.log_trans,\n self.log_emit\n )"
}