call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n// Optimized sigmoid\n__device__ __forceinline__ float fast_sigmoid(float x) {\n return 1.0f / (1.0f + expf(-x));\n}\n\n// Kernel 1: Compute vector norms and concatenate with scalar features\n__global__ void compute_norms_and_combine_kernel(\n const float* __restrict__ scalar_feat, // (B, N, scalar_in)\n const float* __restrict__ vector_feat, // (B, N, vector_in, 3)\n float* __restrict__ combined, // (B, N, scalar_in + vector_in)\n int batch_size, int num_nodes, int scalar_in, int vector_in\n) {\n const int total = batch_size * num_nodes;\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n \n if (idx >= total) return;\n \n const int b = idx / num_nodes;\n const int n = idx % num_nodes;\n \n const int scalar_offset = b * num_nodes * scalar_in + n * scalar_in;\n const int vector_offset = b * num_nodes * vector_in * 3 + n * vector_in * 3;\n const int combined_offset = b * num_nodes * (scalar_in + vector_in) + n * (scalar_in + vector_in);\n \n // Copy scalar features to combined\n for (int s = 0; s < scalar_in; s++) {\n combined[combined_offset + s] = scalar_feat[scalar_offset + s];\n }\n \n // Compute norms and store in combined\n for (int v = 0; v < vector_in; v++) {\n const float vx = vector_feat[vector_offset + v * 3 + 0];\n const float vy = vector_feat[vector_offset + v * 3 + 1];\n const float vz = vector_feat[vector_offset + v * 3 + 2];\n const float norm = sqrtf(vx * vx + vy * vy + vz * vz);\n combined[combined_offset + scalar_in + v] = norm;\n }\n}\n\n// Kernel 2: Scalar linear transform\n__global__ void scalar_linear_kernel(\n const float* __restrict__ combined, // (B, N, combined_dim)\n const float* __restrict__ weight, // (scalar_out, combined_dim)\n const float* __restrict__ bias, // (scalar_out,)\n float* __restrict__ s_out, // (B, N, scalar_out)\n int batch_size, int num_nodes, int combined_dim, int scalar_out\n) {\n const int total = batch_size * num_nodes;\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n \n if (idx >= total) return;\n \n const int b = idx / num_nodes;\n const int n = idx % num_nodes;\n \n const int combined_offset = b * num_nodes * combined_dim + n * combined_dim;\n const int out_offset = b * num_nodes * scalar_out + n * scalar_out;\n \n for (int o = 0; o < scalar_out; o++) {\n float sum = bias[o];\n for (int i = 0; i < combined_dim; i++) {\n sum += combined[combined_offset + i] * weight[o * combined_dim + i];\n }\n s_out[out_offset + o] = sum;\n }\n}\n\n// Kernel 3: Vector linear transform (einsum: oi,bnid->bnod)\n__global__ void vector_linear_kernel(\n const float* __restrict__ vector_feat, // (B, N, vector_in, 3)\n const float* __restrict__ weight, // (vector_out, vector_in)\n float* __restrict__ v_out, // (B, N, vector_out, 3)\n int batch_size, int num_nodes, int vector_in, int vector_out\n) {\n const int total = batch_size * num_nodes;\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n \n if (idx >= total) return;\n \n const int b = idx / num_nodes;\n const int n = idx % num_nodes;\n \n const int in_offset = b * num_nodes * vector_in * 3 + n * vector_in * 3;\n const int out_offset = b * num_nodes * vector_out * 3 + n * vector_out * 3;\n \n for (int o = 0; o < vector_out; o++) {\n float sum_x = 0.0f, sum_y = 0.0f, sum_z = 0.0f;\n for (int i = 0; i < vector_in; i++) {\n const float w = weight[o * vector_in + i];\n sum_x += w * vector_feat[in_offset + i * 3 + 0];\n sum_y += w * vector_feat[in_offset + i * 3 + 1];\n sum_z += w * vector_feat[in_offset + i * 3 + 2];\n }\n v_out[out_offset + o * 3 + 0] = sum_x;\n v_out[out_offset + o * 3 + 1] = sum_y;\n v_out[out_offset + o * 3 + 2] = sum_z;\n }\n}\n\n// Kernel 4: Gating and final output\n__global__ void apply_gate_kernel(\n const float* __restrict__ combined, // (B, N, combined_dim)\n const float* __restrict__ weight, // (vector_out, combined_dim)\n const float* __restrict__ bias, // (vector_out,)\n float* __restrict__ v_out, // (B, N, vector_out, 3)\n int batch_size, int num_nodes, int combined_dim, int vector_out\n) {\n const int total = batch_size * num_nodes;\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n \n if (idx >= total) return;\n \n const int b = idx / num_nodes;\n const int n = idx % num_nodes;\n \n const int combined_offset = b * num_nodes * combined_dim + n * combined_dim;\n const int out_offset = b * num_nodes * vector_out * 3 + n * vector_out * 3;\n \n for (int o = 0; o < vector_out; o++) {\n float sum = bias[o];\n for (int i = 0; i < combined_dim; i++) {\n sum += combined[combined_offset + i] * weight[o * combined_dim + i];\n }\n const float gate = fast_sigmoid(sum);\n \n v_out[out_offset + o * 3 + 0] *= gate;\n v_out[out_offset + o * 3 + 1] *= gate;\n v_out[out_offset + o * 3 + 2] *= gate;\n }\n}\n\n// Host function\nstd::tuple<torch::Tensor, torch::Tensor> se3_forward_cuda(\n torch::Tensor scalar_feat,\n torch::Tensor vector_feat,\n torch::Tensor scalar_weight,\n torch::Tensor scalar_bias,\n torch::Tensor vector_weight,\n torch::Tensor gate_weight,\n torch::Tensor gate_bias\n) {\n const int batch_size = scalar_feat.size(0);\n const int num_nodes = scalar_feat.size(1);\n const int scalar_in = scalar_feat.size(2);\n const int vector_in = vector_feat.size(2);\n const int scalar_out = scalar_weight.size(0);\n const int vector_out = vector_weight.size(0);\n const int combined_dim = scalar_in + vector_in;\n \n const int total = batch_size * num_nodes;\n const int block_size = 256;\n const int grid_size = (total + block_size - 1) / block_size;\n \n // Allocate outputs\n auto combined = torch::empty({batch_size, num_nodes, combined_dim}, scalar_feat.options());\n auto s_out = torch::empty({batch_size, num_nodes, scalar_out}, scalar_feat.options());\n auto v_out = torch::empty({batch_size, num_nodes, vector_out, 3}, scalar_feat.options());\n \n // Kernel 1: Compute norms and combine\n compute_norms_and_combine_kernel<<<grid_size, block_size>>>(\n scalar_feat.data_ptr<float>(),\n vector_feat.data_ptr<float>(),\n combined.data_ptr<float>(),\n batch_size, num_nodes, scalar_in, vector_in\n );\n \n // Kernel 2: Scalar linear\n scalar_linear_kernel<<<grid_size, block_size>>>(\n combined.data_ptr<float>(),\n scalar_weight.data_ptr<float>(),\n scalar_bias.data_ptr<float>(),\n s_out.data_ptr<float>(),\n batch_size, num_nodes, combined_dim, scalar_out\n );\n \n // Kernel 3: Vector linear\n vector_linear_kernel<<<grid_size, block_size>>>(\n vector_feat.data_ptr<float>(),\n vector_weight.data_ptr<float>(),\n v_out.data_ptr<float>(),\n batch_size, num_nodes, vector_in, vector_out\n );\n \n // Kernel 4: Apply gating\n apply_gate_kernel<<<grid_size, block_size>>>(\n combined.data_ptr<float>(),\n gate_weight.data_ptr<float>(),\n gate_bias.data_ptr<float>(),\n v_out.data_ptr<float>(),\n batch_size, num_nodes, combined_dim, vector_out\n );\n \n return std::make_tuple(s_out, v_out);\n}\n\"\"\"\n\ncpp_source = \"\"\"\nstd::tuple<torch::Tensor, torch::Tensor> se3_forward_cuda(\n torch::Tensor scalar_feat,\n torch::Tensor vector_feat,\n torch::Tensor scalar_weight,\n torch::Tensor scalar_bias,\n torch::Tensor vector_weight,\n torch::Tensor gate_weight,\n torch::Tensor gate_bias\n);\n\"\"\"\n\nse3_cuda = load_inline(\n name='se3_cuda_v2',\n cuda_sources=cuda_source,\n cpp_sources=cpp_source,\n functions=['se3_forward_cuda'],\n verbose=False,\n extra_cuda_cflags=['-O3', '--use_fast_math']\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, scalar_in, scalar_out, vector_in, vector_out):\n super().__init__()\n self.scalar_in = scalar_in\n self.scalar_out = scalar_out\n self.vector_in = vector_in\n self.vector_out = vector_out\n \n self.scalar_linear = nn.Linear(scalar_in + vector_in, scalar_out)\n self.vector_weight = nn.Parameter(torch.randn(vector_out, vector_in))\n nn.init.xavier_uniform_(self.vector_weight)\n self.vector_gate = nn.Linear(scalar_in + vector_in, vector_out)\n \n def forward(self, scalar_feat: torch.Tensor, vector_feat: torch.Tensor) -> tuple:\n # Ensure inputs are contiguous float32\n scalar_feat = scalar_feat.contiguous().float()\n vector_feat = vector_feat.contiguous().float()\n \n # Get parameters\n scalar_weight = self.scalar_linear.weight.contiguous().float()\n scalar_bias = self.scalar_linear.bias.contiguous().float()\n vector_weight = self.vector_weight.contiguous().float()\n gate_weight = self.vector_gate.weight.contiguous().float()\n gate_bias = self.vector_gate.bias.contiguous().float()\n \n # Call CUDA kernel\n s_out, v_out = se3_cuda.se3_forward_cuda(\n scalar_feat, vector_feat,\n scalar_weight, scalar_bias,\n vector_weight, gate_weight, gate_bias\n )\n \n return s_out, v_out"
}