call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\nimport os\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cusolverDn.h>\n#include <cublas_v2.h>\n\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)\n\n// Thread block size\n#define BLOCK_SIZE 256\n\n// RBF kernel computation: computes K(X1, X2) \n// X1: (B, N, D), X2: (B, M, D) -> K: (B, N, M)\n// Uses shared memory for X2 tile\n__global__ void rbf_kernel_kernel(\n const float* __restrict__ X1,\n const float* __restrict__ X2,\n float* __restrict__ K,\n const float log_lengthscale,\n const float log_variance,\n int B, int N, int M, int D) {\n \n int b = blockIdx.z;\n int n = blockIdx.y;\n int m = threadIdx.x + blockIdx.x * blockDim.x;\n \n if (b >= B || n >= N || m >= M) return;\n \n float l2 = expf(log_lengthscale) * expf(log_lengthscale);\n float var = expf(log_variance);\n \n float dist_sq = 0.0f;\n for (int d = 0; d < D; d++) {\n float diff = X1[b * N * D + n * D + d] - X2[b * M * D + m * D + d];\n dist_sq += diff * diff;\n }\n \n K[b * N * M + n * M + m] = var * expf(-0.5f * dist_sq / l2);\n}\n\n// Add noise to diagonal\n__global__ void add_noise_diagonal_kernel(\n float* __restrict__ K,\n float noise,\n int B, int N, int M) {\n \n int b = blockIdx.x;\n if (b >= B) return;\n \n for (int i = threadIdx.x; i < N && i < M; i += blockDim.x) {\n K[b * N * M + i * M + i] += noise;\n }\n}\n\n// Add noise diagonal (same size N=N)\n__global__ void add_noise_diagonal_square_kernel(\n float* __restrict__ K,\n float noise,\n int B, int N) {\n \n int b = blockIdx.x;\n int i = threadIdx.x;\n \n if (b >= B || i >= N) return;\n \n K[b * N * N + i * N + i] += noise;\n}\n\n// Compute diagonal of K_mm kernel matrix\n__global__ void kernel_diagonal_kernel(\n const float* __restrict__ X,\n float* __restrict__ diag,\n const float log_lengthscale,\n const float log_variance,\n int B, int N, int D) {\n \n int b = blockIdx.x;\n int n = threadIdx.x;\n \n if (b >= B || n >= N) return;\n \n float var = expf(log_variance);\n // K(x, x) = var * exp(0) = var\n diag[b * N + n] = var;\n}\n\n// Compute mean = K_nm^T @ alpha\n__global__ void compute_mean_kernel(\n const float* __restrict__ K_nm,\n const float* __restrict__ alpha,\n float* __restrict__ mean,\n int B, int N, int M) {\n \n int b = blockIdx.x;\n int m = threadIdx.x;\n \n if (b >= B || m >= M) return;\n \n float sum = 0.0f;\n for (int n = 0; n < N; n++) {\n // K_nm is (N, M), so K_nm^T @ alpha is (M, N) @ (N, 1) = (M, 1)\n // K_nm[b, n, m] = K_nm[b * N * M + n * M + m]\n sum += K_nm[b * N * M + n * M + m] * alpha[b * N + n];\n }\n mean[b * M + m] = sum;\n}\n\n// Compute variance: var = K_mm_diag - (v^2).sum(dim=-2)\n// v is (B, N, M), so v^2 sum over N dimension\n__global__ void compute_variance_kernel(\n const float* __restrict__ K_mm_diag,\n const float* __restrict__ v,\n float* __restrict__ var,\n int B, int N, int M) {\n \n int b = blockIdx.x;\n int m = threadIdx.x;\n \n if (b >= B || m >= M) return;\n \n float sum_sq = 0.0f;\n for (int n = 0; n < N; n++) {\n float v_val = v[b * N * M + n * M + m];\n sum_sq += v_val * v_val;\n }\n \n float result = K_mm_diag[b * M + m] - sum_sq;\n var[b * M + m] = fmaxf(result, 1e-6f);\n}\n\n// Helper class to manage cuSOLVER/cuBLAS handles\nclass SolverHandle {\npublic:\n cusolverDnHandle_t cusolver_handle;\n cublasHandle_t cublas_handle;\n int* d_info;\n \n SolverHandle() {\n cusolverDnCreate(&cusolver_handle);\n cublasCreate(&cublas_handle);\n cudaMalloc(&d_info, sizeof(int));\n }\n \n ~SolverHandle() {\n cusolverDnDestroy(cusolver_handle);\n cublasDestroy(cublas_handle);\n cudaFree(d_info);\n }\n};\n\nstatic SolverHandle& get_solver_handle() {\n static SolverHandle handle;\n return handle;\n}\n\n// Main computation function\nstd::vector<torch::Tensor> gp_forward_cuda(\n torch::Tensor X_train,\n torch::Tensor Y_train,\n torch::Tensor X_test,\n float noise,\n float log_lengthscale,\n float log_variance) {\n \n CHECK_INPUT(X_train);\n CHECK_INPUT(Y_train);\n CHECK_INPUT(X_test);\n \n int B = X_train.size(0);\n int N = X_train.size(1);\n int D = X_train.size(2);\n int M = X_test.size(1);\n \n auto options = torch::TensorOptions()\n .dtype(torch::kFloat32)\n .device(torch::kCUDA)\n .requires_grad(false);\n \n // Allocate output tensors\n torch::Tensor K_nn = torch::empty({B, N, N}, options);\n torch::Tensor K_nm = torch::empty({B, N, M}, options);\n torch::Tensor K_mm_diag = torch::empty({B, M}, options);\n torch::Tensor mean = torch::empty({B, M, 1}, options);\n torch::Tensor var = torch::empty({B, M}, options);\n \n // Get handles\n SolverHandle& handle = get_solver_handle();\n \n // Compute K_nn\n dim3 grid_nn((N + BLOCK_SIZE - 1) / BLOCK_SIZE, N, B);\n dim3 block(BLOCK_SIZE);\n rbf_kernel_kernel<<<grid_nn, block>>>(\n X_train.data_ptr<float>(),\n X_train.data_ptr<float>(),\n K_nn.data_ptr<float>(),\n log_lengthscale, log_variance,\n B, N, N, D);\n \n // Add noise to diagonal of K_nn\n add_noise_diagonal_square_kernel<<<B, N>>>(\n K_nn.data_ptr<float>(), noise, B, N);\n \n // Compute K_nm\n dim3 grid_nm((M + BLOCK_SIZE - 1) / BLOCK_SIZE, N, B);\n rbf_kernel_kernel<<<grid_nm, block>>>(\n X_train.data_ptr<float>(),\n X_test.data_ptr<float>(),\n K_nm.data_ptr<float>(),\n log_lengthscale, log_variance,\n B, N, M, D);\n \n // Compute K_mm diagonal\n kernel_diagonal_kernel<<<B, M>>>(\n X_test.data_ptr<float>(),\n K_mm_diag.data_ptr<float>(),\n log_lengthscale, log_variance,\n B, M, D);\n \n // Cholesky decomposition using cuSOLVER\n torch::Tensor L = K_nn.clone(); // Will be overwritten with L\n \n float** d_A_array = nullptr;\n cudaMalloc(&d_A_array, B * sizeof(float*));\n \n std::vector<float*> h_A_array(B);\n for (int b = 0; b < B; b++) {\n h_A_array[b] = L.data_ptr<float>() + b * N * N;\n }\n cudaMemcpy(d_A_array, h_A_array.data(), B * sizeof(float*), cudaMemcpyHostToDevice);\n \n int lwork = 0;\n cusolverDnSpotrf_bufferSize(\n handle.cusolver_handle,\n CUBLAS_FILL_MODE_LOWER,\n N,\n (float*)L.data_ptr<float>(),\n N,\n &lwork);\n \n float* d_work;\n cudaMalloc(&d_work, lwork * B); // Batched needs more workspace\n \n // Perform batched Cholesky using potrfBatched\n cusolverDnSpotrfBatched(\n handle.cusolver_handle,\n CUBLAS_FILL_MODE_LOWER,\n N,\n d_A_array,\n N,\n handle.d_info,\n B);\n \n // Zero out upper triangle (L should be lower triangular)\n // cuSOLVER returns L in lower triangle, upper is untouched\n // We need to ensure upper is zero for triangular solve\n \n // Solve L @ alpha = Y_train (triangular solve)\n // Y_train is (B, N, 1), so alpha is also (B, N, 1)\n torch::Tensor alpha = Y_train.clone().contiguous();\n \n // Use cublasStrsmBatched for triangular solve\n // L @ alpha = Y_train -> alpha = L^{-1} @ Y_train\n float** d_alpha_array = nullptr;\n cudaMalloc(&d_alpha_array, B * sizeof(float*));\n std::vector<float*> h_alpha_array(B);\n for (int b = 0; b < B; b++) {\n h_alpha_array[b] = alpha.data_ptr<float>() + b * N;\n }\n cudaMemcpy(d_alpha_array, h_alpha_array.data(), B * sizeof(float*), cudaMemcpyHostToDevice);\n \n // alpha = L^{-1} @ Y_train\n // L is lower triangular, so we use CUBLAS_SIDE_LEFT, CUBLAS_FILL_MODE_LOWER\n float one = 1.0f;\n cublasStrsmBatched(\n handle.cublas_handle,\n CUBLAS_SIDE_LEFT,\n CUBLAS_FILL_MODE_LOWER,\n CUBLAS_OP_N,\n CUBLAS_DIAG_NON_UNIT,\n N, 1, // m, n (N x 1 matrix)\n &one,\n d_A_array,\n N,\n d_alpha_array,\n N,\n B);\n \n // Compute mean = K_nm^T @ alpha\n // K_nm is (B, N, M), alpha is (B, N, 1) -> mean is (B, M, 1)\n // This is equivalent to: mean = (K_nm^T @ alpha)\n // Using cublasSgemmBatched\n torch::Tensor mean_flat = torch::empty({B, M}, options);\n \n float** d_Knm_array = nullptr;\n cudaMalloc(&d_Knm_array, B * sizeof(float*));\n std::vector<float*> h_Knm_array(B);\n for (int b = 0; b < B; b++) {\n h_Knm_array[b] = K_nm.data_ptr<float>() + b * N * M;\n }\n cudaMemcpy(d_Knm_array, h_Knm_array.data(), B * sizeof(float*), cudaMemcpyHostToDevice);\n \n float** d_mean_array = nullptr;\n cudaMalloc(&d_mean_array, B * sizeof(float*));\n std::vector<float*> h_mean_array(B);\n for (int b = 0; b < B; b++) {\n h_mean_array[b] = mean_flat.data_ptr<float>() + b * M;\n }\n cudaMemcpy(d_mean_array, h_mean_array.data(), B * sizeof(float*), cudaMemcpyHostToDevice);\n \n // mean = K_nm^T @ alpha\n // K_nm is (N x M), K_nm^T is (M x N), alpha is (N x 1) -> mean is (M x 1)\n // gemm: C = alpha * op(A) * op(B) + beta * C\n // We want: mean = K_nm^T @ alpha\n // A = K_nm^T (M x N), B = alpha (N x 1), C = mean (M x 1)\n float zero = 0.0f;\n cublasSgemmBatched(\n handle.cublas_handle,\n CUBLAS_OP_T, // transpose K_nm\n CUBLAS_OP_N,\n M, 1, N, // m, n, k\n &one,\n d_Knm_array, N, // lda (K_nm is N x M, so lda = N)\n d_alpha_array, N, // ldb (alpha is N x 1, so ldb = N)\n &zero,\n d_mean_array, M, // ldc (mean is M x 1, so ldc = M)\n B);\n \n mean = mean_flat.unsqueeze(-1); // (B, M) -> (B, M, 1)\n \n // Solve L @ v = K_nm for v (triangular solve)\n // v = L^{-1} @ K_nm\n // v is (B, N, M)\n torch::Tensor v = K_nm.clone().contiguous();\n \n float** d_v_array = nullptr;\n cudaMalloc(&d_v_array, B * sizeof(float*));\n std::vector<float*> h_v_array(B);\n for (int b = 0; b < B; b++) {\n h_v_array[b] = v.data_ptr<float>() + b * N * M;\n }\n cudaMemcpy(d_v_array, h_v_array.data(), B * sizeof(float*), cudaMemcpyHostToDevice);\n \n // v = L^{-1} @ K_nm\n // L is (N x N), K_nm is (N x M), v is (N x M)\n cublasStrsmBatched(\n handle.cublas_handle,\n CUBLAS_SIDE_LEFT,\n CUBLAS_FILL_MODE_LOWER,\n CUBLAS_OP_N,\n CUBLAS_DIAG_NON_UNIT,\n N, M, // m, n\n &one,\n d_A_array,\n N,\n d_v_array,\n N,\n B);\n \n // Compute variance: var = K_mm_diag - (v^2).sum(dim=-2)\n // v is (B, N, M), sum over N dimension\n compute_variance_kernel<<<B, M>>>(\n K_mm_diag.data_ptr<float>(),\n v.data_ptr<float>(),\n var.data_ptr<float>(),\n B, N, M);\n \n // Cleanup\n cudaFree(d_A_array);\n cudaFree(d_alpha_array);\n cudaFree(d_Knm_array);\n cudaFree(d_mean_array);\n cudaFree(d_v_array);\n cudaFree(d_work);\n \n return {mean, var};\n}\n\"\"\"\n\ncpp_source = \"\"\"\nstd::vector<torch::Tensor> gp_forward_cuda(\n torch::Tensor X_train,\n torch::Tensor Y_train,\n torch::Tensor X_test,\n float noise,\n float log_lengthscale,\n float log_variance);\n\"\"\"\n\ngp_extension = load_inline(\n name='gp_cuda',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gp_forward_cuda'],\n extra_cuda_cflags=['-lcusolver', '-lcublas'],\n extra_ldflags=['-lcusolver', '-lcublas'],\n verbose=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, input_dim, noise=0.01, lengthscale=1.0, variance=1.0):\n super().__init__()\n self.noise = noise\n self.log_lengthscale = nn.Parameter(torch.tensor(lengthscale).log())\n self.log_variance = nn.Parameter(torch.tensor(variance).log())\n \n def forward(self, X_train: torch.Tensor, Y_train: torch.Tensor, X_test: torch.Tensor) -> tuple:\n # Ensure inputs are contiguous and on GPU\n X_train = X_train.contiguous()\n Y_train = Y_train.contiguous()\n X_test = X_test.contiguous()\n \n if not X_train.is_cuda:\n X_train = X_train.cuda()\n Y_train = Y_train.cuda()\n X_test = X_test.cuda()\n \n mean, var = gp_extension.gp_forward_cuda(\n X_train, Y_train, X_test,\n self.noise,\n self.log_lengthscale.item(),\n self.log_variance.item()\n )\n \n return mean, var"
}