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# Use cuBLAS for matrix operations and fused loss computation\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cublas_v2.h>\n#include <cmath>\n\n// Fused kernel for computing z, KL divergence, and BCE loss\n__global__ void vae_loss_kernel(\n const float* __restrict__ mu,\n const float* __restrict__ logvar,\n const float* __restrict__ eps,\n const float* __restrict__ x_recon,\n const float* __restrict__ x,\n float* __restrict__ z,\n float* __restrict__ elbo_out,\n float* __restrict__ recon_out,\n float* __restrict__ kl_out,\n int batch_size, int latent_dim, int input_dim) {\n \n int batch_idx = blockIdx.x;\n int tid = threadIdx.x;\n \n if (batch_idx >= batch_size) return;\n \n // Pointers to this sample's data\n const float* mu_s = mu + batch_idx * latent_dim;\n const float* logvar_s = logvar + batch_idx * latent_dim;\n const float* eps_s = eps + batch_idx * latent_dim;\n const float* x_recon_s = x_recon + batch_idx * input_dim;\n const float* x_s = x + batch_idx * input_dim;\n float* z_s = z + batch_idx * latent_dim;\n \n // Reparameterization and KL divergence in one pass\n float kl_thread = 0.0f;\n for (int l = tid; l < latent_dim; l += blockDim.x) {\n float m = mu_s[l];\n float lv = logvar_s[l];\n float std_val = expf(0.5f * lv);\n z_s[l] = m + std_val * eps_s[l];\n kl_thread += 1.0f + lv - m * m - expf(lv);\n }\n \n // BCE loss computation\n float recon_thread = 0.0f;\n for (int i = tid; i < input_dim; i += blockDim.x) {\n float xr = x_recon_s[i];\n float xi = x_s[i];\n float log_xr = logf(fmaxf(xr, 1e-7f));\n float log_1_xr = logf(fmaxf(1.0f - xr, 1e-7f));\n recon_thread -= xi * log_xr + (1.0f - xi) * log_1_xr;\n }\n \n // Block reduction\n __shared__ float kl_shared[256];\n __shared__ float recon_shared[256];\n \n kl_shared[tid] = kl_thread;\n recon_shared[tid] = recon_thread;\n __syncthreads();\n \n for (int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (tid < s) {\n kl_shared[tid] += kl_shared[tid + s];\n recon_shared[tid] += recon_shared[tid + s];\n }\n __syncthreads();\n }\n \n if (tid == 0) {\n float kl = -0.5f * kl_shared[0];\n float recon = recon_shared[0];\n elbo_out[batch_idx] = -(recon + kl);\n recon_out[batch_idx] = recon;\n kl_out[batch_idx] = kl;\n }\n}\n\n// Sigmoid activation kernel\n__global__ void sigmoid_kernel(\n const float* __restrict__ x,\n float* __restrict__ y,\n int n) {\n \n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < n) {\n y[idx] = 1.0f / (1.0f + expf(-x[idx]));\n }\n}\n\n// ReLU activation kernel\n__global__ void relu_kernel(\n const float* __restrict__ x,\n float* __restrict__ y,\n int n) {\n \n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < n) {\n y[idx] = fmaxf(x[idx], 0.0f);\n }\n}\n\nstd::vector<torch::Tensor> vae_forward_cuda(\n torch::Tensor x,\n torch::Tensor enc_weight, torch::Tensor enc_bias,\n torch::Tensor mu_weight, torch::Tensor mu_bias,\n torch::Tensor logvar_weight, torch::Tensor logvar_bias,\n torch::Tensor dec1_weight, torch::Tensor dec1_bias,\n torch::Tensor dec2_weight, torch::Tensor dec2_bias,\n torch::Tensor epsilon) {\n \n int batch_size = x.size(0);\n int input_dim = x.size(1);\n int hidden_dim = enc_weight.size(0);\n int latent_dim = mu_weight.size(0);\n \n auto options = torch::TensorOptions().dtype(torch::kFloat32).device(x.device());\n \n // Allocate intermediate tensors\n auto hidden_enc = torch::empty({batch_size, hidden_dim}, options);\n auto mu = torch::empty({batch_size, latent_dim}, options);\n auto logvar = torch::empty({batch_size, latent_dim}, options);\n auto z = torch::empty({batch_size, latent_dim}, options);\n auto hidden_dec = torch::empty({batch_size, hidden_dim}, options);\n auto x_pre_sigmoid = torch::empty({batch_size, input_dim}, options);\n auto x_recon = torch::empty({batch_size, input_dim}, options);\n auto elbo_out = torch::empty({batch_size}, options);\n auto recon_out = torch::empty({batch_size}, options);\n auto kl_out = torch::empty({batch_size}, options);\n \n // Get cuBLAS handle\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n const float alpha = 1.0f;\n const float beta = 0.0f;\n \n // Encoder: hidden = ReLU(x @ enc_weight.T + enc_bias)\n // cublasSgemm: C = alpha * op(A) * op(B) + beta * C\n // We want: hidden = x @ W.T where W is (hidden_dim, input_dim)\n // So: hidden = (W @ x.T).T = x @ W.T\n // Using op(A)=W, op(B)=x.T, so A is W^T no-transpose, B is x transpose\n // Actually easier: use C = A @ B with A=x (batch x in), B=W.T (in x hidden)\n \n // hidden_enc = x @ enc_weight.T + enc_bias\n // enc_weight is (hidden_dim, input_dim), so enc_weight.T is (input_dim, hidden_dim)\n cublasSgemm(handle, \n CUBLAS_OP_N, CUBLAS_OP_T, // op(B)=N, op(A)=T \n hidden_dim, batch_size, input_dim,\n &alpha,\n enc_weight.data_ptr<float>(), input_dim, // A is (hidden x in), op(A)=T gives (in x hidden)\n x.data_ptr<float>(), input_dim, // B is (batch x in), op(B)=N\n &beta,\n hidden_enc.data_ptr<float>(), hidden_dim); // C is (hidden x batch)\n \n // Add bias and apply ReLU (column-major, bias added to each row)\n // hidden_enc is stored as (hidden_dim x batch_size) column-major\n // Add bias to each column\n int relu_threads = 256;\n int relu_blocks = (batch_size * hidden_dim + relu_threads - 1) / relu_threads;\n \n // Need to add bias and ReLU. Let's fuse this.\n auto hidden_enc_2d = hidden_enc.view({hidden_dim, batch_size});\n for (int h = 0; h < hidden_dim; h++) {\n hidden_enc_2d[h].add_(enc_bias[h]);\n }\n relu_kernel<<<relu_blocks, relu_threads>>>(\n hidden_enc.data_ptr<float>(),\n hidden_enc.data_ptr<float>(),\n batch_size * hidden_dim);\n \n // mu = hidden_enc @ mu_weight.T + mu_bias\n // mu_weight is (latent_dim, hidden_dim)\n cublasSgemm(handle,\n CUBLAS_OP_N, CUBLAS_OP_T,\n latent_dim, batch_size, hidden_dim,\n &alpha,\n mu_weight.data_ptr<float>(), hidden_dim,\n hidden_enc.data_ptr<float>(), hidden_dim,\n &beta,\n mu.data_ptr<float>(), latent_dim);\n \n auto mu_2d = mu.view({latent_dim, batch_size});\n for (int l = 0; l < latent_dim; l++) {\n mu_2d[l].add_(mu_bias[l]);\n }\n \n // logvar = hidden_enc @ logvar_weight.T + logvar_bias\n cublasSgemm(handle,\n CUBLAS_OP_N, CUBLAS_OP_T,\n latent_dim, batch_size, hidden_dim,\n &alpha,\n logvar_weight.data_ptr<float>(), hidden_dim,\n hidden_enc.data_ptr<float>(), hidden_dim,\n &beta,\n logvar.data_ptr<float>(), latent_dim);\n \n auto logvar_2d = logvar.view({latent_dim, batch_size});\n for (int l = 0; l < latent_dim; l++) {\n logvar_2d[l].add_(logvar_bias[l]);\n }\n \n // hidden_dec = ReLU(z @ dec1_weight.T + dec1_bias)\n // dec1_weight is (hidden_dim, latent_dim)\n cublasSgemm(handle,\n CUBLAS_OP_N, CUBLAS_OP_T,\n hidden_dim, batch_size, latent_dim,\n &alpha,\n dec1_weight.data_ptr<float>(), latent_dim,\n z.data_ptr<float>(), latent_dim,\n &beta,\n hidden_dec.data_ptr<float>(), hidden_dim);\n \n auto hidden_dec_2d = hidden_dec.view({hidden_dim, batch_size});\n for (int h = 0; h < hidden_dim; h++) {\n hidden_dec_2d[h].add_(dec1_bias[h]);\n }\n relu_kernel<<<relu_blocks, relu_threads>>>(\n hidden_dec.data_ptr<float>(),\n hidden_dec.data_ptr<float>(),\n batch_size * hidden_dim);\n \n // x_pre_sigmoid = hidden_dec @ dec2_weight.T + dec2_bias\n // dec2_weight is (input_dim, hidden_dim)\n cublasSgemm(handle,\n CUBLAS_OP_N, CUBLAS_OP_T,\n input_dim, batch_size, hidden_dim,\n &alpha,\n dec2_weight.data_ptr<float>(), hidden_dim,\n hidden_dec.data_ptr<float>(), hidden_dim,\n &beta,\n x_pre_sigmoid.data_ptr<float>(), input_dim);\n \n auto x_pre_2d = x_pre_sigmoid.view({input_dim, batch_size});\n for (int i = 0; i < input_dim; i++) {\n x_pre_2d[i].add_(dec2_bias[i]);\n }\n \n // Apply sigmoid\n int sig_threads = 256;\n int sig_blocks = (batch_size * input_dim + sig_threads - 1) / sig_threads;\n sigmoid_kernel<<<sig_blocks, sig_threads>>>(\n x_pre_sigmoid.data_ptr<float>(),\n x_recon.data_ptr<float>(),\n batch_size * input_dim);\n \n // Compute z and losses\n int loss_threads = 256;\n int loss_blocks = batch_size;\n vae_loss_kernel<<<loss_blocks, loss_threads>>>(\n mu.data_ptr<float>(),\n logvar.data_ptr<float>(),\n epsilon.data_ptr<float>(),\n x_recon.data_ptr<float>(),\n x.data_ptr<float>(),\n z.data_ptr<float>(),\n elbo_out.data_ptr<float>(),\n recon_out.data_ptr<float>(),\n kl_out.data_ptr<float>(),\n batch_size, latent_dim, input_dim);\n \n return {elbo_out, recon_out, kl_out};\n}\n\"\"\"\n\ncpp_source = \"\"\"\nstd::vector<torch::Tensor> vae_forward_cuda(\n torch::Tensor x,\n torch::Tensor enc_weight, torch::Tensor enc_bias,\n torch::Tensor mu_weight, torch::Tensor mu_bias,\n torch::Tensor logvar_weight, torch::Tensor logvar_bias,\n torch::Tensor dec1_weight, torch::Tensor dec1_bias,\n torch::Tensor dec2_weight, torch::Tensor dec2_bias,\n torch::Tensor epsilon);\n\"\"\"\n\nvae_extension = load_inline(\n name='vae_extension',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['vae_forward_cuda'],\n extra_cuda_cflags=['-lcublas'],\n libraries=['cublas'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self, input_dim, hidden_dim, latent_dim):\n super().__init__()\n self.encoder = nn.Sequential(\n nn.Linear(input_dim, hidden_dim),\n nn.ReLU(),\n )\n self.fc_mu = nn.Linear(hidden_dim, latent_dim)\n self.fc_logvar = nn.Linear(hidden_dim, latent_dim)\n\n self.decoder = nn.Sequential(\n nn.Linear(latent_dim, hidden_dim),\n nn.ReLU(),\n nn.Linear(hidden_dim, input_dim),\n nn.Sigmoid(),\n )\n \n self.input_dim = input_dim\n self.hidden_dim = hidden_dim\n self.latent_dim = latent_dim\n\n def forward(self, x: torch.Tensor) -> tuple:\n batch_size = x.size(0)\n epsilon = torch.randn(batch_size, self.latent_dim, device=x.device, dtype=x.dtype)\n \n results = vae_extension.vae_forward_cuda(\n x.contiguous(),\n self.encoder[0].weight.contiguous(),\n self.encoder[0].bias.contiguous(),\n self.fc_mu.weight.contiguous(),\n self.fc_mu.bias.contiguous(),\n self.fc_logvar.weight.contiguous(),\n self.fc_logvar.bias.contiguous(),\n self.decoder[0].weight.contiguous(),\n self.decoder[0].bias.contiguous(),\n self.decoder[2].weight.contiguous(),\n self.decoder[2].bias.contiguous(),\n epsilon)\n \n return results[0], results[1], results[2]"
}