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\n// Predict kernel: one thread per particle\ntemplate <typename scalar_t>\n__global__ void predict_kernel(\n const scalar_t* __restrict__ particles,\n const scalar_t* __restrict__ A,\n scalar_t* __restrict__ pred,\n int B, int P, int D,\n unsigned long long seed) {\n \n int b = blockIdx.x;\n int p = threadIdx.x + blockIdx.y * blockDim.x;\n \n if (b >= B || p >= P) return;\n \n // Initialize random state for this particle\n curandState state;\n curand_init(seed, b * P + p, 0, &state);\n \n // Compute x @ A.T for this particle (all dimensions)\n for (int d = 0; d < D; d++) {\n scalar_t val = 0;\n for (int k = 0; k < D; k++) {\n val += particles[b * P * D + p * D + k] * A[d * D + k];\n }\n // Add process noise (0.1 * randn)\n val += scalar_t(0.1) * curand_normal(&state);\n pred[b * P * D + p * D + d] = val;\n }\n}\n\n// Weight kernel: compute log likelihood\ntemplate <typename scalar_t>\n__global__ void weight_lik_kernel(\n const scalar_t* __restrict__ pred,\n const scalar_t* __restrict__ C,\n const scalar_t* __restrict__ observation,\n const scalar_t* __restrict__ log_weights,\n scalar_t* __restrict__ new_log_w,\n int B, int P, int D, int obs_dim,\n scalar_t obs_noise_sq) {\n \n int b = blockIdx.x;\n int p = threadIdx.x + blockIdx.y * blockDim.x;\n \n if (b >= B || p >= P) return;\n \n // Compute obs_pred = pred @ C.T\n scalar_t log_lik = 0;\n for (int o = 0; o < obs_dim; o++) {\n scalar_t obs_pred_o = 0;\n for (int d = 0; d < D; d++) {\n obs_pred_o += pred[b * P * D + p * D + d] * C[o * D + d];\n }\n scalar_t diff = observation[b * obs_dim + o] - obs_pred_o;\n log_lik += -0.5 * diff * diff / obs_noise_sq;\n }\n new_log_w[b * P + p] = log_weights[b * P + p] + log_lik;\n}\n\n// Normalize log weights (logsumexp)\ntemplate <typename scalar_t>\n__global__ void normalize_kernel(\n scalar_t* __restrict__ new_log_w,\n int B, int P) {\n \n int b = blockIdx.x;\n if (b >= B) return;\n \n // Shared memory for reduction\n extern __shared__ char smem_ptr[];\n scalar_t* smem = reinterpret_cast<scalar_t*>(smem_ptr);\n \n int tid = threadIdx.x;\n int p = tid;\n \n // Find max\n scalar_t my_val = (p < P) ? new_log_w[b * P + p] : -INFINITY;\n \n // First, reduce within each warp\n for (int offset = 16; offset > 0; offset >>= 1) {\n scalar_t other = __shfl_xor_sync(0xffffffff, my_val, offset);\n my_val = max(my_val, other);\n }\n \n // Store warp results to shared memory\n int lane_id = tid % 32;\n int warp_id = tid / 32;\n int num_warps = (blockDim.x + 31) / 32;\n \n if (lane_id == 0) {\n smem[warp_id] = my_val;\n }\n __syncthreads();\n \n // Reduce across warps (first warp only)\n if (warp_id == 0) {\n scalar_t val = (tid < num_warps) ? smem[tid] : -INFINITY;\n for (int offset = 16; offset > 0; offset >>= 1) {\n scalar_t other = __shfl_xor_sync(0xffffffff, val, offset);\n val = max(val, other);\n }\n if (tid == 0) {\n smem[0] = val; // max_log_w\n }\n }\n __syncthreads();\n \n scalar_t max_log_w = smem[0];\n __syncthreads();\n \n // Compute exp(val - max) for each\n scalar_t my_exp = (p < P) ? exp(new_log_w[b * P + p] - max_log_w) : 0;\n \n // Reduce within each warp\n for (int offset = 16; offset > 0; offset >>= 1) {\n my_exp += __shfl_xor_sync(0xffffffff, my_exp, offset);\n }\n \n if (lane_id == 0) {\n smem[warp_id] = my_exp;\n }\n __syncthreads();\n \n // Reduce across warps\n if (warp_id == 0) {\n scalar_t val = (tid < num_warps) ? smem[tid] : 0;\n for (int offset = 16; offset > 0; offset >>= 1) {\n val += __shfl_xor_sync(0xffffffff, val, offset);\n }\n if (tid == 0) {\n smem[0] = val; // sum_exp\n }\n }\n __syncthreads();\n \n scalar_t sum_exp = smem[0];\n scalar_t log_sum = max_log_w + log(sum_exp + 1e-30);\n \n if (p < P) {\n new_log_w[b * P + p] -= log_sum;\n }\n}\n\n// Resample kernel: systematic resampling with cumulative sum\ntemplate <typename scalar_t>\n__global__ void resample_kernel(\n const scalar_t* __restrict__ pred,\n const scalar_t* __restrict__ weights,\n scalar_t* __restrict__ resampled,\n int B, int P, int D,\n unsigned long long seed) {\n \n int b = blockIdx.x;\n int p = threadIdx.x + blockIdx.y * blockDim.x;\n \n if (b >= B || p >= P) return;\n \n // Shared memory for cumulative sum (only first block of each batch uses it)\n extern __shared__ char smem_ptr[];\n scalar_t* cumsum = reinterpret_cast<scalar_t*>(smem_ptr);\n \n int tid = threadIdx.x;\n \n // Load weights into shared memory\n scalar_t my_w = (p < P) ? weights[b * P + p] : 0;\n cumsum[tid] = my_w;\n __syncthreads();\n \n // Parallel prefix sum (inclusive scan) using Blelloch/work-efficient algorithm\n // Since we're doing it within a block, we need to handle P threads\n // For simplicity, use sequential addressing (O(n^2) but small D)\n // Actually let's use parallel prefix sum\n int n = min(P, (int)blockDim.x);\n \n // Up-sweep (reduce)\n for (int offset = 1; offset < n; offset *= 2) {\n int idx = tid + offset;\n scalar_t val = 0;\n if (idx < n && tid % (2*offset) == 0) {\n val = cumsum[tid + offset];\n }\n __syncthreads();\n if (idx < n && tid % (2*offset) == 0) {\n cumsum[tid + offset] += cumsum[tid];\n }\n __syncthreads();\n }\n \n // Actually let me use a simpler approach: each thread computes cumsum up to its position\n // But we need to handle multiple blocks...\n // Let me rethink this\n \n // For systematic resampling, we need:\n // 1. Cumulative sum of weights\n // 2. Binary search for each (p + u_offset) / P\n \n // Since P might be > 1024, we can't fit all in shared memory\n // Let's use a different approach: compute cumsum in global memory\n \n // Actually, for this problem size (P=1024), we can use one block per batch\n // and do parallel scan in shared memory\n}\n\n// Alternative resample kernel using global memory for cumsum\ntemplate <typename scalar_t>\n__global__ void cumsum_kernel(\n const scalar_t* __restrict__ weights,\n scalar_t* __restrict__ cumsum,\n int B, int P) {\n \n int b = blockIdx.x;\n int p = threadIdx.x + blockIdx.y * blockDim.x;\n \n if (b >= B || p >= P) return;\n \n // Sequential cumsum per batch (inefficient but correct)\n if (p == 0) {\n cumsum[b * P + 0] = weights[b * P + 0];\n for (int i = 1; i < P; i++) {\n cumsum[b * P + i] = cumsum[b * P + i-1] + weights[b * P + i];\n }\n }\n}\n\ntemplate <typename scalar_t>\n__global__ void resample_gather_kernel(\n const scalar_t* __restrict__ pred,\n const scalar_t* __restrict__ cumsum,\n scalar_t* __restrict__ resampled,\n scalar_t* __restrict__ u_offsets,\n int B, int P, int D,\n unsigned long long seed) {\n \n int b = blockIdx.x;\n int p = threadIdx.x + blockIdx.y * blockDim.x;\n \n if (b >= B || p >= P) return;\n \n // Generate stratified uniform\n scalar_t u_offset = u_offsets[b];\n scalar_t u = (scalar_t(p) + u_offset) / scalar_t(P);\n \n // Binary search for index\n int lo = 0, hi = P - 1;\n while (lo < hi) {\n int mid = (lo + hi) / 2;\n scalar_t mid_val = cumsum[b * P + mid];\n if (mid_val < u) {\n lo = mid + 1;\n } else {\n hi = mid;\n }\n }\n int idx = min(lo, P - 1);\n \n // Gather from pred\n for (int d = 0; d < D; d++) {\n resampled[b * P * D + p * D + d] = pred[b * P * D + idx * D + d];\n }\n}\n\ntemplate <typename scalar_t>\n__global__ void init_u_offsets_kernel(\n scalar_t* __restrict__ u_offsets,\n int B,\n unsigned long long seed) {\n \n int b = blockIdx.x;\n if (b >= B) return;\n \n curandState state;\n curand_init(seed, b, 0, &state);\n u_offsets[b] = curand_uniform(&state);\n}\n\nstd::tuple<torch::Tensor, torch::Tensor> particle_filter_cuda(\n torch::Tensor particles,\n torch::Tensor log_weights,\n torch::Tensor observation,\n torch::Tensor A,\n torch::Tensor C,\n int state_dim,\n int obs_dim,\n int num_particles,\n double process_noise,\n double obs_noise) {\n \n int B = particles.size(0);\n int P = num_particles;\n int D = state_dim;\n \n auto options = particles.options();\n auto pred = torch::empty({B, P, D}, options);\n auto new_log_w = torch::empty({B, P}, options);\n auto resampled = torch::empty({B, P, D}, options);\n auto cumsum = torch::empty({B, P}, options);\n auto u_offsets = torch::empty({B}, options);\n \n // Random seeds\n unsigned long long seed1 = std::chrono::system_clock::now().time_since_epoch().count();\n unsigned long long seed3 = seed1 + 67890;\n \n int threads_per_block = 256;\n int blocks_y = (P + threads_per_block - 1) / threads_per_block;\n \n // Predict kernel\n dim3 block(threads_per_block);\n dim3 grid(B, blocks_y);\n \n AT_DISPATCH_FLOATING_TYPES(particles.scalar_type(), \"predict_kernel\", ([&] {\n predict_kernel<scalar_t><<<grid, block>>>(\n particles.data_ptr<scalar_t>(),\n A.data_ptr<scalar_t>(),\n pred.data_ptr<scalar_t>(),\n B, P, D, seed1);\n }));\n \n // Weight kernel (compute log likelihood)\n AT_DISPATCH_FLOATING_TYPES(particles.scalar_type(), \"weight_lik_kernel\", ([&] {\n weight_lik_kernel<scalar_t><<<grid, block>>>(\n pred.data_ptr<scalar_t>(),\n C.data_ptr<scalar_t>(),\n observation.data_ptr<scalar_t>(),\n log_weights.data_ptr<scalar_t>(),\n new_log_w.data_ptr<scalar_t>(),\n B, P, D, obs_dim, scalar_t(obs_noise * obs_noise));\n }));\n \n // Normalize (logsumexp)\n int smem_size = ((threads_per_block + 31) / 32) * sizeof(float); // For warp reductions\n AT_DISPATCH_FLOATING_TYPES(particles.scalar_type(), \"normalize_kernel\", ([&] {\n int smem_size_typed = ((threads_per_block + 31) / 32) * sizeof(scalar_t);\n normalize_kernel<scalar_t><<<B, threads_per_block, smem_size_typed>>>(\n new_log_w.data_ptr<scalar_t>(),\n B, P);\n }));\n \n // Convert to weights for resampling\n auto weights = torch::exp(new_log_w);\n \n // Compute cumulative sum\n AT_DISPATCH_FLOATING_TYPES(particles.scalar_type(), \"cumsum_kernel\", ([&] {\n cumsum_kernel<scalar_t><<<B, 1>>>(\n weights.data_ptr<scalar_t>(),\n cumsum.data_ptr<scalar_t>(),\n B, P);\n }));\n \n // Initialize random offsets\n AT_DISPATCH_FLOATING_TYPES(particles.scalar_type(), \"init_u_offsets_kernel\", ([&] {\n init_u_offsets_kernel<scalar_t><<<B, 1>>>(\n u_offsets.data_ptr<scalar_t>(),\n B, seed3);\n }));\n \n // Resample gather\n AT_DISPATCH_FLOATING_TYPES(particles.scalar_type(), \"resample_gather_kernel\", ([&] {\n resample_gather_kernel<scalar_t><<<grid, block>>>(\n pred.data_ptr<scalar_t>(),\n cumsum.data_ptr<scalar_t>(),\n resampled.data_ptr<scalar_t>(),\n u_offsets.data_ptr<scalar_t>(),\n B, P, D, seed3);\n }));\n \n auto uniform_log_w = torch::full({B, P}, -std::log(double(P)), options);\n \n return std::make_tuple(resampled, uniform_log_w);\n}\n\"\"\"\n\ncpp_source = \"\"\"\nstd::tuple<torch::Tensor, torch::Tensor> particle_filter_cuda(\n torch::Tensor particles,\n torch::Tensor log_weights,\n torch::Tensor observation,\n torch::Tensor A,\n torch::Tensor C,\n int state_dim,\n int obs_dim,\n int num_particles,\n double process_noise,\n double obs_noise);\n\"\"\"\n\nparticle_filter_module = load_inline(\n name=\"particle_filter_module_v3\",\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=[\"particle_filter_cuda\"],\n extra_cuda_cflags=[\"-lcurand\"],\n verbose=True,\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self, state_dim, obs_dim, num_particles):\n super().__init__()\n self.state_dim = state_dim\n self.obs_dim = obs_dim\n self.num_particles = num_particles\n self.register_buffer(\"A\", 0.9 * torch.eye(state_dim))\n self.register_buffer(\"C\", torch.randn(obs_dim, state_dim) * 0.5)\n self.process_noise = 0.1\n self.obs_noise = 0.5\n\n def forward(self, particles, log_weights, observation):\n resampled, uniform_log_w = particle_filter_module.particle_filter_cuda(\n particles, log_weights, observation, self.A, self.C,\n self.state_dim, self.obs_dim, self.num_particles,\n self.process_noise, self.obs_noise\n )\n return resampled, uniform_log_w"
}