call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r'''\n#include <torch/extension.h>\n\ntorch::Tensor rbf_forward_cuda(torch::Tensor distances, int64_t num_rbf, double cutoff);\n\ntorch::Tensor rbf_forward(torch::Tensor distances, int64_t num_rbf, double cutoff) {\n return rbf_forward_cuda(distances, num_rbf, cutoff);\n}\n'''\n\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <vector>\n\n__global__ void rbf64_kernel(const float* __restrict__ dists,\n float* __restrict__ out,\n int64_t M,\n float cutoff,\n float coeff) {\n // 4 distances per block, 64 threads/features per distance = 256 threads\n constexpr int R = 64;\n constexpr int DPB = 4;\n __shared__ float s_d[DPB];\n __shared__ float s_env[DPB];\n int tid = threadIdx.x;\n int g = tid >> 6; // 0..3\n int r = tid & 63; // 0..63\n int64_t m = (int64_t)blockIdx.x * DPB + g;\n if (m >= M) return;\n if (r == 0) {\n float d = dists[m];\n s_d[g] = d;\n float env = 0.0f;\n if (d < cutoff) {\n env = 0.5f * (__cosf(3.14159265358979323846f * d / cutoff) + 1.0f);\n }\n s_env[g] = env;\n }\n __syncthreads();\n float d = s_d[g];\n float off = (cutoff * (float)r) * (1.0f / 63.0f);\n float x = d - off;\n float val = __expf(coeff * x * x) * s_env[g];\n out[m * R + r] = val;\n}\n\n__global__ void rbf_generic_kernel(const float* __restrict__ dists,\n float* __restrict__ out,\n int64_t M,\n int R,\n float cutoff,\n float coeff,\n float inv_rminus1) {\n int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;\n int64_t total = M * (int64_t)R;\n if (idx >= total) return;\n int r = (int)(idx % R);\n int64_t m = idx / R;\n float d = dists[m];\n float env = 0.0f;\n if (d < cutoff) {\n env = 0.5f * (__cosf(3.14159265358979323846f * d / cutoff) + 1.0f);\n }\n float off = cutoff * (float)r * inv_rminus1;\n float x = d - off;\n out[idx] = __expf(coeff * x * x) * env;\n}\n\ntorch::Tensor rbf_forward_cuda(torch::Tensor distances, int64_t num_rbf, double cutoff_d) {\n TORCH_CHECK(distances.is_cuda(), \"distances must be CUDA\");\n TORCH_CHECK(distances.scalar_type() == torch::kFloat32, \"distances must be float32\");\n auto dist = distances.contiguous();\n std::vector<int64_t> sizes = dist.sizes().vec();\n sizes.push_back(num_rbf);\n auto out = torch::empty(sizes, dist.options());\n int64_t M = dist.numel();\n if (M == 0) return out;\n float cutoff = (float)cutoff_d;\n float step = cutoff / (float)(num_rbf - 1);\n float coeff = -0.5f / (step * step);\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if (num_rbf == 64) {\n dim3 block(256);\n dim3 grid((unsigned int)((M + 3) / 4));\n rbf64_kernel<<<grid, block, 0, stream>>>(dist.data_ptr<float>(), out.data_ptr<float>(), M, cutoff, coeff);\n } else {\n int threads = 256;\n int64_t total = M * num_rbf;\n int blocks = (int)((total + threads - 1) / threads);\n float inv = 1.0f / (float)(num_rbf - 1);\n rbf_generic_kernel<<<blocks, threads, 0, stream>>>(dist.data_ptr<float>(), out.data_ptr<float>(), M, (int)num_rbf, cutoff, coeff, inv);\n }\n return out;\n}\n'''\n\n_rbf_ext = load_inline(\n name='rbf_ext_v1',\n cpp_sources=_cpp_src,\n cuda_sources=_cuda_src,\n functions=['rbf_forward'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n extra_cflags=['-O3'],\n verbose=False,\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, num_rbf=64, cutoff=5.0):\n super().__init__()\n self.num_rbf = num_rbf\n self.cutoff = cutoff\n offsets = torch.linspace(0.0, cutoff, num_rbf)\n self.register_buffer('offsets', offsets)\n self.coeff = -0.5 / (offsets[1] - offsets[0]).item() ** 2\n\n def forward(self, distances: torch.Tensor) -> torch.Tensor:\n return _rbf_ext.rbf_forward(distances, self.num_rbf, float(self.cutoff))\n"
}