call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r'''\n#include <torch/extension.h>\n#include <vector>\n\nstd::vector<torch::Tensor> smc_forward_cuda(torch::Tensor particles,\n torch::Tensor log_weights,\n torch::Tensor observation,\n torch::Tensor C,\n torch::Tensor noise,\n torch::Tensor u0);\n\nstd::vector<torch::Tensor> smc_forward(torch::Tensor particles,\n torch::Tensor log_weights,\n torch::Tensor observation,\n torch::Tensor C,\n torch::Tensor noise,\n torch::Tensor u0) {\n return smc_forward_cuda(particles, log_weights, observation, C, noise, u0);\n}\n'''\n\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <vector>\n#include <cmath>\n\n__global__ void compute_pred_logw_kernel(const float* __restrict__ particles,\n const float* __restrict__ log_weights,\n const float* __restrict__ obs,\n const float* __restrict__ C,\n const float* __restrict__ noise,\n float* __restrict__ pred,\n float* __restrict__ raw_logw,\n int P) {\n int b = blockIdx.x;\n int tid = threadIdx.x;\n float c00 = C[0], c01 = C[1], c02 = C[2], c03 = C[3];\n float c10 = C[4], c11 = C[5], c12 = C[6], c13 = C[7];\n float o0 = obs[b * 2 + 0];\n float o1 = obs[b * 2 + 1];\n int bp = b * P;\n for (int i = tid; i < P; i += blockDim.x) {\n int off4 = (bp + i) * 4;\n float x0 = 0.9f * particles[off4 + 0] + 0.1f * noise[off4 + 0];\n float x1 = 0.9f * particles[off4 + 1] + 0.1f * noise[off4 + 1];\n float x2 = 0.9f * particles[off4 + 2] + 0.1f * noise[off4 + 2];\n float x3 = 0.9f * particles[off4 + 3] + 0.1f * noise[off4 + 3];\n pred[off4 + 0] = x0;\n pred[off4 + 1] = x1;\n pred[off4 + 2] = x2;\n pred[off4 + 3] = x3;\n float y0 = c00*x0 + c01*x1 + c02*x2 + c03*x3;\n float y1 = c10*x0 + c11*x1 + c12*x2 + c13*x3;\n float d0 = o0 - y0;\n float d1 = o1 - y1;\n raw_logw[bp + i] = log_weights[bp + i] - 2.0f * (d0*d0 + d1*d1);\n }\n}\n\n__global__ void scan_resample_kernel(const float* __restrict__ pred,\n const float* __restrict__ raw_logw,\n const float* __restrict__ u0,\n float* __restrict__ out_particles,\n float* __restrict__ out_logw,\n int P) {\n __shared__ float s[1024];\n int b = blockIdx.x;\n int tid = threadIdx.x;\n int bp = b * P;\n\n float v = (tid < P) ? raw_logw[bp + tid] : -CUDART_INF_F;\n s[tid] = v;\n __syncthreads();\n // max reduction (blockDim is 1024)\n for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {\n if (tid < stride) {\n float other = s[tid + stride];\n if (other > s[tid]) s[tid] = other;\n }\n __syncthreads();\n }\n float m = s[0];\n\n float w = (tid < P) ? expf(v - m) : 0.0f;\n s[tid] = w;\n __syncthreads();\n for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {\n if (tid < stride) s[tid] += s[tid + stride];\n __syncthreads();\n }\n float inv_sum = 1.0f / s[0];\n float wi = (tid < P) ? w * inv_sum : 0.0f;\n s[tid] = wi;\n __syncthreads();\n\n // Inclusive prefix sum of normalized weights.\n for (int offset = 1; offset < 1024; offset <<= 1) {\n float add = 0.0f;\n if (tid >= offset) add = s[tid - offset];\n __syncthreads();\n s[tid] += add;\n __syncthreads();\n }\n\n if (tid < P) {\n float target = ((float)tid + u0[b]) / (float)P;\n int lo = 0;\n int hi = P;\n while (lo < hi) {\n int mid = (lo + hi) >> 1;\n float cm = s[mid];\n if (cm < target) lo = mid + 1;\n else hi = mid;\n }\n int idx = lo;\n if (idx >= P) idx = P - 1;\n int out4 = (bp + tid) * 4;\n int in4 = (bp + idx) * 4;\n out_particles[out4 + 0] = pred[in4 + 0];\n out_particles[out4 + 1] = pred[in4 + 1];\n out_particles[out4 + 2] = pred[in4 + 2];\n out_particles[out4 + 3] = pred[in4 + 3];\n out_logw[bp + tid] = -logf((float)P);\n }\n}\n\nstd::vector<torch::Tensor> smc_forward_cuda(torch::Tensor particles,\n torch::Tensor log_weights,\n torch::Tensor observation,\n torch::Tensor C,\n torch::Tensor noise,\n torch::Tensor u0) {\n TORCH_CHECK(particles.is_cuda(), \"particles must be CUDA\");\n TORCH_CHECK(particles.scalar_type() == torch::kFloat32, \"float32 only\");\n TORCH_CHECK(particles.dim() == 3 && particles.size(2) == 4, \"expected D=4\");\n TORCH_CHECK(observation.size(1) == 2, \"expected obs_dim=2\");\n int B = particles.size(0);\n int P = particles.size(1);\n TORCH_CHECK(P <= 1024, \"P must be <=1024\");\n auto pred = torch::empty_like(particles);\n auto raw = torch::empty_like(log_weights);\n auto out = torch::empty_like(particles);\n auto out_logw = torch::empty_like(log_weights);\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n compute_pred_logw_kernel<<<B, 256, 0, stream>>>(\n particles.data_ptr<float>(), log_weights.data_ptr<float>(), observation.data_ptr<float>(),\n C.data_ptr<float>(), noise.data_ptr<float>(), pred.data_ptr<float>(), raw.data_ptr<float>(), P);\n scan_resample_kernel<<<B, 1024, 0, stream>>>(\n pred.data_ptr<float>(), raw.data_ptr<float>(), u0.data_ptr<float>(),\n out.data_ptr<float>(), out_logw.data_ptr<float>(), P);\n return {out, out_logw};\n}\n'''\n\n_ext = load_inline(name='smc_ext_v1', cpp_sources=_cpp_src, cuda_sources=_cuda_src,\n functions=['smc_forward'], extra_cuda_cflags=['-O3'], verbose=False)\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: torch.Tensor, log_weights: torch.Tensor, observation: torch.Tensor) -> tuple:\n noise = torch.randn_like(particles)\n u0 = torch.rand((particles.shape[0],), device=particles.device, dtype=particles.dtype)\n return tuple(_ext.smc_forward(particles, log_weights, observation, self.C, noise, u0))\n"
}