call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n# Load and compile the CUDA kernel\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n// SiLU activation: x * sigmoid(x)\n__device__ __forceinline__ float silu(float x) {\n return x * (1.0f / (1.0f + expf(-x)));\n}\n\n// Warp-level matrix-vector multiply: each warp computes one output row\n// input: [IN_DIM], weight: [OUT_DIM * IN_DIM], bias: [OUT_DIM]\n// Each thread reads a portion of input and weight, then reduces within warp\ntemplate<int IN_DIM, int OUT_DIM>\n__device__ void warp_linear_silu(\n const float* __restrict__ input,\n const float* __restrict__ weight,\n const float* __restrict__ bias,\n float* output,\n int out_idx, // which output row this warp computes\n int lane_id\n) {\n float sum = 0.0f;\n \n // Strided access to input - each thread reads multiple elements\n for (int i = lane_id; i < IN_DIM; i += 32) {\n sum += input[i] * weight[out_idx * IN_DIM + i];\n }\n \n // Reduce within warp\n #pragma unroll\n for (int offset = 16; offset > 0; offset >>= 1) {\n sum += __shfl_down_sync(0xffffffff, sum, offset);\n }\n \n // Lane 0 adds bias and writes output\n if (lane_id == 0) {\n output[out_idx] = silu(sum + bias[out_idx]);\n }\n}\n\ntemplate<int IN_DIM, int OUT_DIM>\n__device__ void warp_linear(\n const float* __restrict__ input,\n const float* __restrict__ weight,\n const float* __restrict__ bias,\n float* output,\n int out_idx,\n int lane_id,\n bool apply_silu\n) {\n float sum = 0.0f;\n \n for (int i = lane_id; i < IN_DIM; i += 32) {\n sum += input[i] * weight[out_idx * IN_DIM + i];\n }\n \n #pragma unroll\n for (int offset = 16; offset > 0; offset >>= 1) {\n sum += __shfl_down_sync(0xffffffff, sum, offset);\n }\n \n if (lane_id == 0) {\n sum += bias[out_idx];\n output[out_idx] = apply_silu ? silu(sum) : sum;\n }\n}\n\n// Main kernel: fused ODE integration\n// Grid: (batch_size * num_warps_per_sample) blocks\n// Block: 32 threads (one warp)\n__global__ void fused_ode_kernel(\n const float* __restrict__ x0,\n const float* __restrict__ w1,\n const float* __restrict__ b1,\n const float* __restrict__ w2,\n const float* __restrict__ b2,\n const float* __restrict__ w3,\n const float* __restrict__ b3,\n float* __restrict__ x_out,\n int dim,\n int hidden_dim,\n int num_steps,\n int batch_size\n) {\n int lane_id = threadIdx.x;\n int warp_id = blockIdx.x;\n int batch_idx = warp_id / (hidden_dim / 32); // Each sample uses hidden_dim/32 warps\n int local_warp_id = warp_id % (hidden_dim / 32);\n \n if (batch_idx >= batch_size || lane_id >= 32) return;\n \n // Shared memory for this sample's data\n // We'll use global memory for x and intermediates, shared for tiling if needed\n // Actually, let's have each warp compute 32 output elements\n \n // Each warp computes 32 consecutive output elements\n int out_start = local_warp_id * 32;\n \n // Pointer to this sample's x in global memory\n // We need shared memory to hold x for the integration\n // Use a different approach: each block processes one batch sample\n // with multiple warps cooperating\n \n // Simplified: process one batch per block with multiple warps\n}\n\n// Alternative: one block per sample, multiple warps per block\n__global__ void fused_ode_kernel_v2(\n const float* __restrict__ x0,\n const float* __restrict__ w1,\n const float* __restrict__ b1,\n const float* __restrict__ w2,\n const float* __restrict__ b2,\n const float* __restrict__ w3,\n const float* __restrict__ b3,\n float* __restrict__ x_out,\n int dim,\n int hidden_dim,\n int num_steps,\n int batch_size\n) {\n int batch_idx = blockIdx.x;\n int lane_id = threadIdx.x & 31;\n int warp_id = threadIdx.x >> 5;\n int warps_per_block = blockDim.x >> 5;\n \n if (batch_idx >= batch_size) return;\n \n // Shared memory\n extern __shared__ char smem[];\n float* s_x = reinterpret_cast<float*>(smem); // dim\n float* s_input = s_x + dim; // dim + 1\n float* s_hidden1 = s_input + dim + 1; // hidden_dim\n float* s_hidden2 = s_hidden1 + hidden_dim; // hidden_dim\n \n // Initialize x from x0 (cooperative load)\n for (int i = threadIdx.x; i < dim; i += blockDim.x) {\n s_x[i] = x0[batch_idx * dim + i];\n }\n __syncthreads();\n \n float dt = 1.0f / float(num_steps);\n \n // Euler integration\n for (int step = 0; step < num_steps; step++) {\n float t = (float(step) + 0.5f) / float(num_steps);\n \n // Build input [x, t]\n for (int i = threadIdx.x; i < dim; i += blockDim.x) {\n s_input[i] = s_x[i];\n }\n if (threadIdx.x == 0) {\n s_input[dim] = t;\n }\n __syncthreads();\n \n // Layer 1: (dim+1) -> hidden_dim with SiLU\n // Each warp computes multiple outputs\n int outputs_per_warp = (hidden_dim + warps_per_block - 1) / warps_per_block;\n int my_out_start = warp_id * outputs_per_warp;\n int my_out_end = min(my_out_start + outputs_per_warp, hidden_dim);\n \n for (int out_idx = my_out_start + lane_id; out_idx < my_out_end; out_idx += 32) {\n float sum = b1[out_idx];\n for (int i = 0; i < dim + 1; i++) {\n sum += s_input[i] * w1[out_idx * (dim + 1) + i];\n }\n s_hidden1[out_idx] = silu(sum);\n }\n __syncthreads();\n \n // Layer 2: hidden_dim -> hidden_dim with SiLU\n outputs_per_warp = (hidden_dim + warps_per_block - 1) / warps_per_block;\n my_out_start = warp_id * outputs_per_warp;\n my_out_end = min(my_out_start + outputs_per_warp, hidden_dim);\n \n for (int out_idx = my_out_start + lane_id; out_idx < my_out_end; out_idx += 32) {\n float sum = b2[out_idx];\n for (int i = 0; i < hidden_dim; i++) {\n sum += s_hidden1[i] * w2[out_idx * hidden_dim + i];\n }\n s_hidden2[out_idx] = silu(sum);\n }\n __syncthreads();\n \n // Layer 3: hidden_dim -> dim (no activation)\n outputs_per_warp = (dim + warps_per_block - 1) / warps_per_block;\n my_out_start = warp_id * outputs_per_warp;\n my_out_end = min(my_out_start + outputs_per_warp, dim);\n \n for (int out_idx = my_out_start + lane_id; out_idx < my_out_end; out_idx += 32) {\n float sum = b3[out_idx];\n for (int i = 0; i < hidden_dim; i++) {\n sum += s_hidden2[i] * w3[out_idx * hidden_dim + i];\n }\n s_x[out_idx] = s_x[out_idx] + dt * sum;\n }\n __syncthreads();\n }\n \n // Write output\n for (int i = threadIdx.x; i < dim; i += blockDim.x) {\n x_out[batch_idx * dim + i] = s_x[i];\n }\n}\n\n// Host function\ntorch::Tensor fused_ode_forward(\n torch::Tensor x0,\n torch::Tensor w1,\n torch::Tensor b1,\n torch::Tensor w2,\n torch::Tensor b2,\n torch::Tensor w3,\n torch::Tensor b3,\n int64_t num_steps\n) {\n TORCH_CHECK(x0.is_cuda(), \"x0 must be CUDA tensor\");\n TORCH_CHECK(x0.dim() == 2, \"x0 must be 2D\");\n \n int batch_size = x0.size(0);\n int dim = x0.size(1);\n int hidden_dim = w1.size(0);\n \n auto x_out = torch::empty_like(x0);\n \n // Shared memory: x + input + hidden1 + hidden2\n int smem_size = (dim + dim + 1 + hidden_dim * 2) * sizeof(float);\n \n // Use 256 or 512 threads per block for better occupancy\n int threads = 256;\n if (hidden_dim <= 128) threads = 128;\n if (hidden_dim > 256) threads = 512;\n threads = min(threads, 1024);\n \n fused_ode_kernel_v2<<<batch_size, threads, smem_size>>>(\n x0.data_ptr<float>(),\n w1.data_ptr<float>(),\n b1.data_ptr<float>(),\n w2.data_ptr<float>(),\n b2.data_ptr<float>(),\n w3.data_ptr<float>(),\n b3.data_ptr<float>(),\n x_out.data_ptr<float>(),\n dim,\n hidden_dim,\n num_steps,\n batch_size\n );\n \n return x_out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor fused_ode_forward(\n torch::Tensor x0,\n torch::Tensor w1,\n torch::Tensor b1,\n torch::Tensor w2,\n torch::Tensor b2,\n torch::Tensor w3,\n torch::Tensor b3,\n int64_t num_steps\n);\n\"\"\"\n\n# Load the extension\nfused_ode_ext = load_inline(\n name='fused_ode_ext_v2',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['fused_ode_forward'],\n extra_cuda_cflags=['-O3', '-use_fast_math'],\n verbose=False\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self, dim: int, num_steps: int):\n super().__init__()\n assert dim >= 1 and num_steps >= 1\n self.dim = dim\n self.num_steps = num_steps\n self.register_buffer(\"dt\", torch.tensor(1.0 / float(num_steps)))\n \n # Velocity network: same architecture as reference\n self.velocity = nn.Sequential(\n nn.Linear(dim + 1, dim * 4),\n nn.SiLU(),\n nn.Linear(dim * 4, dim * 4),\n nn.SiLU(),\n nn.Linear(dim * 4, dim),\n )\n \n def forward(self, x0: torch.Tensor) -> torch.Tensor:\n B, d = x0.shape\n assert d == self.dim\n \n # Sequential structure:\n # [0]: Linear(dim+1, dim*4)\n # [1]: SiLU\n # [2]: Linear(dim*4, dim*4)\n # [3]: SiLU\n # [4]: Linear(dim*4, dim)\n layers = list(self.velocity.children())\n w1 = layers[0].weight\n b1 = layers[0].bias\n w2 = layers[2].weight\n b2 = layers[2].bias\n w3 = layers[4].weight\n b3 = layers[4].bias\n \n # Ensure contiguous\n w1 = w1.contiguous()\n b1 = b1.contiguous()\n w2 = w2.contiguous()\n b2 = b2.contiguous()\n w3 = w3.contiguous()\n b3 = b3.contiguous()\n x0 = x0.contiguous()\n \n return fused_ode_ext.fused_ode_forward(\n x0, w1, b1, w2, b2, w3, b3, self.num_steps\n )"
}