call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\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 <vector>\n#include <cmath>\n\n// GELU approximation (fast)\n__device__ __forceinline__ float gelu_fast(float x) {\n // tanh approximation\n return 0.5f * x * (1.0f + tanhf(0.7978845608f * (x + 0.044715f * x * x * x)));\n}\n\n// Exact GELU\n__device__ __forceinline__ float gelu_exact(float x) {\n return 0.5f * x * (1.0f + erff(x * 0.7071067811865475f));\n}\n\n// Warp-level reduce sum\n__device__ float warp_reduce_sum(float val) {\n #pragma unroll\n for (int offset = 16; offset > 0; offset /= 2) {\n val += __shfl_down_sync(0xffffffff, val, offset);\n }\n return val;\n}\n\n// Block-level reduce sum using shared memory\n__device__ float block_reduce_sum(float val, float* shared) {\n int lane = threadIdx.x % 32;\n int wid = threadIdx.x / 32;\n \n val = warp_reduce_sum(val);\n \n if (lane == 0) {\n shared[wid] = val;\n }\n __syncthreads();\n \n val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f;\n if (wid == 0) {\n val = warp_reduce_sum(val);\n }\n return val;\n}\n\n// Fused LayerNorm + GELU kernel\n// Normalizes over [channels, seq_len] for each batch element\n__global__ void layernorm_gelu_kernel(\n const float* __restrict__ input,\n float* __restrict__ output,\n int batch_size,\n int channels,\n int seq_len\n) {\n int b = blockIdx.x;\n int tid = threadIdx.x;\n int total = channels * seq_len;\n \n // Load weight/bias from global memory (assume weight=1, bias=0 for simplicity)\n // In the actual model, these are learned parameters\n \n extern __shared__ float shared_mem[];\n float* shared_sum = shared_mem;\n float* shared_sq_sum = shared_mem + blockDim.x / 32;\n \n float local_sum = 0.0f;\n float local_sq_sum = 0.0f;\n \n // Each thread processes multiple elements\n for (int i = tid; i < total; i += blockDim.x) {\n int c = i / seq_len;\n int s = i % seq_len;\n int idx = b * channels * seq_len + c * seq_len + s;\n float val = input[idx];\n local_sum += val;\n local_sq_sum += val * val;\n }\n \n float sum = block_reduce_sum(local_sum, shared_sum);\n __syncthreads();\n float sq_sum = block_reduce_sum(local_sq_sum, shared_sq_sum);\n \n float mean = sum / total;\n float var = sq_sum / total - mean * mean;\n float inv_std = rsqrtf(var + 1e-5f);\n \n // Apply normalization and GELU\n for (int i = tid; i < total; i += blockDim.x) {\n int c = i / seq_len;\n int s = i % seq_len;\n int idx = b * channels * seq_len + c * seq_len + s;\n float val = (input[idx] - mean) * inv_std;\n output[idx] = gelu_fast(val);\n }\n}\n\n// Dilated Conv1d kernel with residual\n// Grid: (batch_size, channels_out, seq_len) or tiled\n__global__ void dilated_conv1d_residual_kernel(\n const float* __restrict__ input,\n const float* __restrict__ weight,\n const float* __restrict__ bias,\n const float* __restrict__ residual,\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) {\n int b = blockIdx.x;\n int out_c = blockIdx.y;\n int out_s = blockIdx.z * blockDim.x + threadIdx.x;\n \n if (out_s >= seq_len) return;\n \n // Compute the starting position in the input sequence\n // For Conv1d with dilation: input position = out_s * 1 + k * dilation - padding\n float sum = bias[out_c];\n \n int in_s_center = out_s - padding; // Center position in input\n \n for (int k = 0; k < kernel_size; k++) {\n int in_s = in_s_center + k * dilation;\n if (in_s >= 0 && in_s < seq_len) {\n for (int in_c = 0; in_c < channels; in_c++) {\n float w = weight[out_c * channels * kernel_size + in_c * kernel_size + k];\n float v = input[b * channels * seq_len + in_c * seq_len + in_s];\n sum += w * v;\n }\n }\n }\n \n // Add residual\n sum += residual[b * channels * seq_len + out_c * seq_len + out_s];\n \n output[b * channels * seq_len + out_c * seq_len + out_s] = sum;\n}\n\n// Optimized dilated Conv1d with shared memory for weights\n__global__ void dilated_conv1d_residual_opt_kernel(\n const float* __restrict__ input,\n const float* __restrict__ weight,\n const float* __restrict__ bias,\n const float* __restrict__ residual,\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) {\n int b = blockIdx.x;\n int out_c = blockIdx.y;\n \n // Shared memory for weight slice [in_c, k]\n extern __shared__ float shared_weights[];\n \n // Each block processes multiple output positions\n // Load weights for this output channel into shared memory\n for (int i = threadIdx.x; i < channels * kernel_size; i += blockDim.x) {\n shared_weights[i] = weight[out_c * channels * kernel_size + i];\n }\n __syncthreads();\n \n float bias_val = bias[out_c];\n \n for (int out_s = threadIdx.x; out_s < seq_len; out_s += blockDim.x) {\n float sum = bias_val;\n int in_s_center = out_s - padding;\n \n for (int k = 0; k < kernel_size; k++) {\n int in_s = in_s_center + k * dilation;\n if (in_s >= 0 && in_s < seq_len) {\n for (int in_c = 0; in_c < channels; in_c++) {\n float w = shared_weights[in_c * kernel_size + k];\n float v = input[b * channels * seq_len + in_c * seq_len + in_s];\n sum += w * v;\n }\n }\n }\n \n sum += residual[b * channels * seq_len + out_c * seq_len + out_s];\n output[b * channels * seq_len + out_c * seq_len + out_s] = sum;\n }\n}\n\n// Host functions\nvoid launch_layernorm_gelu(\n const float* input,\n float* output,\n int batch_size,\n int channels,\n int seq_len,\n cudaStream_t stream\n) {\n int total = channels * seq_len;\n int threads = min(1024, total);\n int shared_mem = (threads / 32 + threads / 32) * sizeof(float);\n \n layernorm_gelu_kernel<<<batch_size, threads, shared_mem, stream>>>(\n input, output, batch_size, channels, seq_len\n );\n}\n\nvoid launch_dilated_conv1d_residual(\n const float* input,\n const float* weight,\n const float* bias,\n const float* residual,\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 cudaStream_t stream\n) {\n // Use optimized kernel with shared memory for weights\n int threads = 256;\n int shared_mem = channels * kernel_size * sizeof(float);\n \n dim3 grid(batch_size, channels);\n \n dilated_conv1d_residual_opt_kernel<<<grid, threads, shared_mem, stream>>>(\n input, weight, bias, residual, output,\n batch_size, channels, seq_len, kernel_size, dilation, padding\n );\n}\n\"\"\"\n\ncpp_source = \"\"\"\nvoid launch_layernorm_gelu(const float* input, float* output, int batch_size, int channels, int seq_len, cudaStream_t stream);\nvoid launch_dilated_conv1d_residual(const float* input, const float* weight, const float* bias, const float* residual, float* output, int batch_size, int channels, int seq_len, int kernel_size, int dilation, int padding, cudaStream_t stream);\n\ntorch::Tensor layernorm_gelu_cuda(torch::Tensor input) {\n auto output = torch::empty_like(input);\n int batch_size = input.size(0);\n int channels = input.size(1);\n int seq_len = input.size(2);\n \n launch_layernorm_gelu(\n input.data_ptr<float>(),\n output.data_ptr<float>(),\n batch_size, channels, seq_len,\n at::cuda::getCurrentCUDAStream()\n );\n return output;\n}\n\ntorch::Tensor dilated_conv1d_residual_cuda(\n torch::Tensor input,\n torch::Tensor weight,\n torch::Tensor bias,\n torch::Tensor residual\n) {\n auto output = torch::empty_like(input);\n int batch_size = input.size(0);\n int channels = input.size(1);\n int seq_len = input.size(2);\n int kernel_size = weight.size(2);\n int dilation = 1; // This will be set per layer\n int padding = dilation * (kernel_size / 2);\n \n launch_dilated_conv1d_residual(\n input.data_ptr<float>(),\n weight.data_ptr<float>(),\n bias.data_ptr<float>(),\n residual.data_ptr<float>(),\n output.data_ptr<float>(),\n batch_size, channels, seq_len,\n kernel_size, dilation, padding,\n at::cuda::getCurrentCUDAStream()\n );\n return output;\n}\n\"\"\"\n\n# Load the CUDA extension\ncuda_extension = load_inline(\n name='dilated_conv_cuda',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['layernorm_gelu_cuda', 'dilated_conv1d_residual_cuda'],\n verbose=True,\n extra_cuda_cflags=['-O3']\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self, channels, seq_len, kernel_size, max_dilation_power):\n super(ModelNew, self).__init__()\n self.channels = channels\n self.seq_len = seq_len\n self.kernel_size = kernel_size\n self.max_dilation_power = max_dilation_power\n \n # Store norms and convs for parameter access\n self.norms = nn.ModuleList()\n self.convs = nn.ModuleList()\n self.activation = nn.GELU()\n \n for i in range(max_dilation_power + 1):\n dilation = 2 ** i\n self.norms.append(nn.LayerNorm([channels, seq_len]))\n self.convs.append(nn.Conv1d(channels, channels, kernel_size, padding=dilation, dilation=dilation))\n \n def forward(self, x):\n residual = x\n for i, (norm, conv) in enumerate(zip(self.norms, self.convs)):\n # Use PyTorch for now to verify correctness\n x = norm(x)\n x = self.activation(x)\n x = conv(x)\n x = x + residual\n return x"
}