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\n\n# CUDA kernel code\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <curand_kernel.h>\n\ntemplate <typename scalar_t>\n__global__ void cvae_forward_kernel(\n const scalar_t* __restrict__ x,\n const scalar_t* __restrict__ y,\n const scalar_t* __restrict__ enc_weight,\n const scalar_t* __restrict__ enc_bias,\n const scalar_t* __restrict__ enc_mu_weight,\n const scalar_t* __restrict__ enc_mu_bias,\n const scalar_t* __restrict__ enc_logvar_weight,\n const scalar_t* __restrict__ enc_logvar_bias,\n const scalar_t* __restrict__ dec_weight,\n const scalar_t* __restrict__ dec_bias,\n const scalar_t* __restrict__ dec_out_weight,\n const scalar_t* __restrict__ dec_out_bias,\n scalar_t* __restrict__ loss_out,\n scalar_t* __restrict__ recon_loss_out,\n scalar_t* __restrict__ kl_loss_out,\n int batch_size,\n int input_dim,\n int label_dim,\n int hidden_dim,\n int latent_dim\n) {\n // Each thread handles one sample in the batch\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= batch_size) return;\n \n // Shared memory for weights would help, but let's start simple\n // Encoder: concatenate x and y\n int enc_in_dim = input_dim + label_dim;\n \n // Allocate local arrays for intermediate results\n extern __shared__ char shared_mem[];\n scalar_t* enc_hidden = reinterpret_cast<scalar_t*>(shared_mem);\n scalar_t* mu = enc_hidden + blockDim.x * hidden_dim;\n scalar_t* logvar = mu + blockDim.x * latent_dim;\n scalar_t* z = logvar + blockDim.x * latent_dim;\n scalar_t* dec_hidden = z + blockDim.x * latent_dim;\n scalar_t* x_recon = dec_hidden + blockDim.x * hidden_dim;\n \n int tid = threadIdx.x;\n scalar_t* my_enc_hidden = enc_hidden + tid * hidden_dim;\n scalar_t* my_mu = mu + tid * latent_dim;\n scalar_t* my_logvar = logvar + tid * latent_dim;\n scalar_t* my_z = z + tid * latent_dim;\n scalar_t* my_dec_hidden = dec_hidden + tid * hidden_dim;\n scalar_t* my_x_recon = x_recon + tid * input_dim;\n \n // Encoder linear layer: h = ReLU(W * [x,y] + b)\n for (int h = 0; h < hidden_dim; h++) {\n scalar_t sum = enc_bias[h];\n for (int i = 0; i < input_dim; i++) {\n sum += x[idx * input_dim + i] * enc_weight[h * enc_in_dim + i];\n }\n for (int i = 0; i < label_dim; i++) {\n sum += y[idx * label_dim + i] * enc_weight[h * enc_in_dim + input_dim + i];\n }\n my_enc_hidden[h] = sum > 0 ? sum : 0; // ReLU\n }\n \n // Compute mu\n for (int l = 0; l < latent_dim; l++) {\n scalar_t sum = enc_mu_bias[l];\n for (int h = 0; h < hidden_dim; h++) {\n sum += my_enc_hidden[h] * enc_mu_weight[l * hidden_dim + h];\n }\n my_mu[l] = sum;\n }\n \n // Compute logvar\n for (int l = 0; l < latent_dim; l++) {\n scalar_t sum = enc_logvar_bias[l];\n for (int h = 0; h < hidden_dim; h++) {\n sum += my_enc_hidden[h] * enc_logvar_weight[l * hidden_dim + h];\n }\n my_logvar[l] = sum;\n }\n \n // Sample z = mu + std * noise\n curandState state;\n curand_init(idx, 0, 0, &state);\n \n for (int l = 0; l < latent_dim; l++) {\n scalar_t std_val = exp(0.5f * my_logvar[l]);\n scalar_t noise = curand_normal(&state);\n my_z[l] = my_mu[l] + std_val * noise;\n }\n \n // Decoder: concatenate z and y\n int dec_in_dim = latent_dim + label_dim;\n \n // Decoder linear layer: h = ReLU(W * [z,y] + b)\n for (int h = 0; h < hidden_dim; h++) {\n scalar_t sum = dec_bias[h];\n for (int l = 0; l < latent_dim; l++) {\n sum += my_z[l] * dec_weight[h * dec_in_dim + l];\n }\n for (int i = 0; i < label_dim; i++) {\n sum += y[idx * label_dim + i] * dec_weight[h * dec_in_dim + latent_dim + i];\n }\n my_dec_hidden[h] = sum > 0 ? sum : 0; // ReLU\n }\n \n // Decoder output layer: x_recon = W * h + b\n for (int i = 0; i < input_dim; i++) {\n scalar_t sum = dec_out_bias[i];\n for (int h = 0; h < hidden_dim; h++) {\n sum += my_dec_hidden[h] * dec_out_weight[i * hidden_dim + h];\n }\n my_x_recon[i] = sum;\n }\n \n // Compute recon_loss = MSE(x_recon, x)\n scalar_t recon_loss = 0;\n for (int i = 0; i < input_dim; i++) {\n scalar_t diff = my_x_recon[i] - x[idx * input_dim + i];\n recon_loss += diff * diff;\n }\n \n // Compute kl_loss = -0.5 * sum(1 + logvar - mu^2 - exp(logvar))\n scalar_t kl_loss = 0;\n for (int l = 0; l < latent_dim; l++) {\n kl_loss += 1 + my_logvar[l] - my_mu[l] * my_mu[l] - exp(my_logvar[l]);\n }\n kl_loss = -0.5f * kl_loss;\n \n // Store results\n recon_loss_out[idx] = recon_loss;\n kl_loss_out[idx] = kl_loss;\n loss_out[idx] = recon_loss + kl_loss;\n}\n\n// Reduction kernel to sum losses\ntemplate <typename scalar_t>\n__global__ void sum_reduce_kernel(\n const scalar_t* __restrict__ input,\n scalar_t* __restrict__ output,\n int n\n) {\n __shared__ scalar_t sdata[256];\n int tid = threadIdx.x;\n int i = blockIdx.x * blockDim.x + threadIdx.x;\n \n sdata[tid] = (i < n) ? input[i] : 0;\n __syncthreads();\n \n for (int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (tid < s) {\n sdata[tid] += sdata[tid + s];\n }\n __syncthreads();\n }\n \n if (tid == 0) {\n atomicAdd(output, sdata[0]);\n }\n}\n\ntorch::Tensor cvae_forward(\n torch::Tensor x,\n torch::Tensor y,\n torch::Tensor enc_weight,\n torch::Tensor enc_bias,\n torch::Tensor enc_mu_weight,\n torch::Tensor enc_mu_bias,\n torch::Tensor enc_logvar_weight,\n torch::Tensor enc_logvar_bias,\n torch::Tensor dec_weight,\n torch::Tensor dec_bias,\n torch::Tensor dec_out_weight,\n torch::Tensor dec_out_bias\n) {\n int batch_size = x.size(0);\n int input_dim = x.size(1);\n int label_dim = y.size(1);\n int hidden_dim = enc_weight.size(0);\n int latent_dim = enc_mu_weight.size(0);\n \n auto options = x.options();\n auto loss_out = torch::zeros({batch_size}, options);\n auto recon_loss_out = torch::zeros({batch_size}, options);\n auto kl_loss_out = torch::zeros({batch_size}, options);\n \n int threads = 256;\n int blocks = (batch_size + threads - 1) / threads;\n \n // Shared memory for intermediate results\n int shared_mem_size = threads * (hidden_dim + latent_dim + latent_dim + hidden_dim + input_dim) * sizeof(float);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"cvae_forward_kernel\", ([&] {\n cvae_forward_kernel<scalar_t><<<blocks, threads, shared_mem_size>>>(\n x.data_ptr<scalar_t>(),\n y.data_ptr<scalar_t>(),\n enc_weight.data_ptr<scalar_t>(),\n enc_bias.data_ptr<scalar_t>(),\n enc_mu_weight.data_ptr<scalar_t>(),\n enc_mu_bias.data_ptr<scalar_t>(),\n enc_logvar_weight.data_ptr<scalar_t>(),\n enc_logvar_bias.data_ptr<scalar_t>(),\n dec_weight.data_ptr<scalar_t>(),\n dec_bias.data_ptr<scalar_t>(),\n dec_out_weight.data_ptr<scalar_t>(),\n dec_out_bias.data_ptr<scalar_t>(),\n loss_out.data_ptr<scalar_t>(),\n recon_loss_out.data_ptr<scalar_t>(),\n kl_loss_out.data_ptr<scalar_t>(),\n batch_size,\n input_dim,\n label_dim,\n hidden_dim,\n latent_dim\n );\n }));\n \n // Sum the losses\n auto total_recon = torch::zeros({1}, options);\n auto total_kl = torch::zeros({1}, options);\n auto total_loss = torch::zeros({1}, options);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"sum_reduce_kernel\", ([&] {\n sum_reduce_kernel<scalar_t><<<blocks, threads>>>(\n recon_loss_out.data_ptr<scalar_t>(),\n total_recon.data_ptr<scalar_t>(),\n batch_size\n );\n sum_reduce_kernel<scalar_t><<<blocks, threads>>>(\n kl_loss_out.data_ptr<scalar_t>(),\n total_kl.data_ptr<scalar_t>(),\n batch_size\n );\n sum_reduce_kernel<scalar_t><<<blocks, threads>>>(\n loss_out.data_ptr<scalar_t>(),\n total_loss.data_ptr<scalar_t>(),\n batch_size\n );\n }));\n \n // Divide by batch size\n total_recon.div_(batch_size);\n total_kl.div_(batch_size);\n total_loss.div_(batch_size);\n \n return torch::stack({total_loss, total_recon, total_kl});\n}\n\"\"\"\n\n# Load the CUDA extension\ncvae_ops = load_inline(\n name=\"cvae_ops\",\n cpp_sources=\"\",\n cuda_sources=cuda_source,\n functions=[\"cvae_forward\"],\n extra_cuda_cflags=[\"-lcurand\"],\n verbose=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, input_dim, label_dim, hidden_dim, latent_dim):\n super().__init__()\n # Encoder: q(z | x, y)\n self.enc = nn.Sequential(\n nn.Linear(input_dim + label_dim, hidden_dim),\n nn.ReLU(),\n )\n self.enc_mu = nn.Linear(hidden_dim, latent_dim)\n self.enc_logvar = nn.Linear(hidden_dim, latent_dim)\n\n # Decoder: p(x | z, y)\n self.dec = nn.Sequential(\n nn.Linear(latent_dim + label_dim, hidden_dim),\n nn.ReLU(),\n nn.Linear(hidden_dim, input_dim),\n )\n\n def forward(self, x: torch.Tensor, y: torch.Tensor) -> tuple:\n # Get weights and biases\n enc_weight = self.enc[0].weight # shape: (hidden_dim, input_dim + label_dim)\n enc_bias = self.enc[0].bias # shape: (hidden_dim,)\n enc_mu_weight = self.enc_mu.weight # shape: (latent_dim, hidden_dim)\n enc_mu_bias = self.enc_mu.bias # shape: (latent_dim,)\n enc_logvar_weight = self.enc_logvar.weight # shape: (latent_dim, hidden_dim)\n enc_logvar_bias = self.enc_logvar.bias # shape: (latent_dim,)\n \n # Decoder layers\n dec_weight = self.dec[0].weight # shape: (hidden_dim, latent_dim + label_dim)\n dec_bias = self.dec[0].bias # shape: (hidden_dim,)\n dec_out_weight = self.dec[2].weight # shape: (input_dim, hidden_dim)\n dec_out_bias = self.dec[2].bias # shape: (input_dim,)\n \n # Call CUDA kernel\n result = cvae_ops.cvae_forward(\n x, y,\n enc_weight, enc_bias,\n enc_mu_weight, enc_mu_bias,\n enc_logvar_weight, enc_logvar_bias,\n dec_weight, dec_bias,\n dec_out_weight, dec_out_bias\n )\n \n loss = result[0].squeeze()\n recon_loss = result[1].squeeze()\n kl_loss = result[2].squeeze()\n \n return loss, recon_loss, kl_loss"
}