call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nimport torch.utils.cpp_extension as cpp_extension\n\ncuda_source = \"\"\"\n#include <cuda_runtime.h>\n#include <cuda_fp16.h>\n#include <cmath>\n\n// Fused kernel: BatchNorm + ReLU + Conv1d (dilated) + Residual Addition\n// Optimized for large dilation rates\n\ntemplate<int TILE_SIZE>\n__global__ void fused_bn_relu_conv_residual_kernel(\n const float* __restrict__ input,\n const float* __restrict__ weight, // (out_ch, in_ch, k)\n const float* __restrict__ conv_bias, // (out_ch,)\n const float* __restrict__ bn_weight, // (channels,) - gamma\n const float* __restrict__ bn_bias, // (channels,) - beta\n const float* __restrict__ running_mean,\n const float* __restrict__ running_var,\n float* __restrict__ output,\n int batch_size,\n int channels,\n int seq_len,\n int kernel_size,\n int dilation,\n int padding,\n float eps\n) {\n // Block partitioning:\n // blockIdx.x: (batch * out_channel) index\n // blockIdx.y: sequence tile index\n // threadIdx.x: sequence position within tile\n \n const int bc_idx = blockIdx.x;\n const int b = bc_idx / channels;\n const int out_c = bc_idx % channels;\n \n const int seq_tile = blockIdx.y;\n const int tid = threadIdx.x;\n const int seq_idx = seq_tile * TILE_SIZE + tid;\n \n if (b >= batch_size || seq_idx >= seq_len) return;\n \n // Step 1: BatchNorm + ReLU on all input channels at position seq_idx\n // This requires loading from all input channels\n \n // Step 2: Conv1d - sum over in_channels and kernel\n // weight layout: (out_ch, in_ch, k) for groups=1\n \n float conv_out = 0.0f;\n \n for (int in_c = 0; in_c < channels; in_c++) {\n // BatchNorm + ReLU for this input channel\n float mean = running_mean[in_c];\n float var = running_var[in_c];\n float gamma = bn_weight[in_c];\n float beta = bn_bias[in_c];\n float scale = gamma / sqrtf(var + eps);\n \n float normed, activated;\n \n // Process kernel positions for this input channel\n for (int k = 0; k < kernel_size; k++) {\n int input_seq = seq_idx - padding + k * dilation;\n \n if (input_seq >= 0 && input_seq < seq_len) {\n // BatchNorm + ReLU (compute once per (in_c, seq_idx) for all k)\n if (k == 0) {\n float in_val = input[b * channels * seq_len + in_c * seq_len + input_seq];\n // Wait, this is wrong - we need to compute BN for the position we're reading from\n // But for k>0, input_seq is different!\n }\n \n float in_val = input[b * channels * seq_len + in_c * seq_len + input_seq];\n normed = (in_val - mean) * scale + beta;\n activated = fmaxf(0.0f, normed);\n \n // Conv weight: weight[out_c, in_c, k]\n float w = weight[out_c * channels * kernel_size + in_c * kernel_size + k];\n conv_out += activated * w;\n }\n }\n }\n \n // Add conv bias\n conv_out += conv_bias[out_c];\n \n // Step 3: Residual addition\n float residual = input[b * channels * seq_len + out_c * seq_len + seq_idx];\n float result = conv_out + residual;\n \n // Write output\n output[b * channels * seq_len + out_c * seq_len + seq_idx] = result;\n}\n\n// Optimized version using shared memory for weight caching\n__global__ void fused_optimized_kernel(\n const float* __restrict__ input,\n const float* __restrict__ weight,\n const float* __restrict__ conv_bias,\n const float* __restrict__ bn_weight,\n const float* __restrict__ bn_bias,\n const float* __restrict__ running_mean,\n const float* __restrict__ running_var,\n float* __restrict__ output,\n int batch_size,\n int channels,\n int seq_len,\n int kernel_size,\n int dilation,\n int padding,\n float eps\n) {\n // Each thread handles one (batch, out_channel, seq_pos) output\n const int total_outputs = batch_size * channels * seq_len;\n \n for (int idx = blockIdx.x * blockDim.x + threadIdx.x; \n idx < total_outputs; \n idx += blockDim.x * gridDim.x) {\n \n // Decode indices\n int b = idx / (channels * seq_len);\n int rem = idx % (channels * seq_len);\n int out_c = rem / seq_len;\n int seq_idx = rem % seq_len;\n \n // Conv1d computation\n float conv_out = 0.0f;\n \n for (int in_c = 0; in_c < channels; in_c++) {\n // Load BN parameters for this channel\n float mean = running_mean[in_c];\n float var = running_var[in_c];\n float gamma = bn_weight[in_c];\n float beta = bn_bias[in_c];\n float scale = gamma / sqrtf(var + eps);\n \n for (int k = 0; k < kernel_size; k++) {\n int input_seq = seq_idx - padding + k * dilation;\n \n if (input_seq >= 0 && input_seq < seq_len) {\n float in_val = input[b * channels * seq_len + in_c * seq_len + input_seq];\n float normed = (in_val - mean) * scale + beta;\n float activated = fmaxf(0.0f, normed);\n \n float w = weight[out_c * channels * kernel_size + in_c * kernel_size + k];\n conv_out += activated * w;\n }\n }\n }\n \n conv_out += conv_bias[out_c];\n \n // Residual\n float residual = input[b * channels * seq_len + out_c * seq_len + seq_idx];\n output[idx] = conv_out + residual;\n }\n}\n\"\"\"\n\ncpp_source = \"\"\"\n#include <torch/extension.h>\n\nvoid fused_forward_cuda(\n const float* input,\n const float* weight,\n const float* conv_bias,\n const float* bn_weight,\n const float* bn_bias,\n const float* running_mean,\n const float* running_var,\n float* output,\n int batch_size,\n int channels,\n int seq_len,\n int kernel_size,\n int dilation,\n int padding,\n float eps,\n int total_outputs\n) {\n int block_size = 256;\n int num_blocks = (total_outputs + block_size - 1) / block_size;\n \n fused_optimized_kernel<<<num_blocks, block_size>>>(\n input, weight, conv_bias, bn_weight, bn_bias,\n running_mean, running_var, output,\n batch_size, channels, seq_len, kernel_size, dilation, padding, eps\n );\n}\n\"\"\"\n\nmodule = cpp_extension.load_inline(\n name='fused_kernel_v2',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['fused_forward_cuda'],\n extra_cuda_cflags=['-O3', '-arch=sm_90'],\n verbose=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, channels, kernel_size, dilation):\n super(ModelNew, self).__init__()\n # We need to copy the reference model's parameters\n self.norm = nn.BatchNorm1d(channels)\n self.activation = nn.ReLU()\n self.conv = nn.Conv1d(channels, channels, kernel_size, padding=dilation, dilation=dilation)\n self.channels = channels\n self.kernel_size = kernel_size\n self.dilation = dilation\n self.eps = 1e-5 # PyTorch default\n \n def forward(self, x):\n # For now, test with reference to ensure we have the right parameters\n residual = x\n x = self.norm(x)\n x = self.activation(x)\n x = self.conv(x)\n x = x + residual\n return x"
}