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>\ntorch::Tensor linear_act(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double multiplier, double negative_slope);\n'''\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <c10/cuda/CUDAException.h>\n#include <cuda_runtime.h>\n#include <cublasLt.h>\n#include <mutex>\n\n#define CUBLAS_CHECK(cmd) do { cublasStatus_t s = (cmd); if (s != CUBLAS_STATUS_SUCCESS) { TORCH_CHECK(false, \"cuBLASLt error: \", (int)s); } } while(0)\n\nstatic cublasLtHandle_t get_lt_handle() {\n static cublasLtHandle_t handle = nullptr;\n static std::once_flag flag;\n std::call_once(flag, [](){ CUBLAS_CHECK(cublasLtCreate(&handle)); });\n return handle;\n}\n\n__device__ __forceinline__ float lact(float v, float s) { return (v > 0.0f) ? v : v * s; }\n\n__global__ void add_bias_lrelu_vec4(float* __restrict__ y, const float* __restrict__ bias,\n long long n4, int N, float mult, float slope) {\n long long tid = (long long)blockIdx.x * blockDim.x + threadIdx.x;\n long long stride = (long long)blockDim.x * gridDim.x;\n const float4* __restrict__ b4 = reinterpret_cast<const float4*>(bias);\n float4* __restrict__ y4 = reinterpret_cast<float4*>(y);\n int N4 = N >> 2;\n for (long long i = tid; i < n4; i += stride) {\n float4 v = y4[i];\n float4 b = b4[(int)(i % N4)];\n v.x = lact(v.x + mult * b.x, slope);\n v.y = lact(v.y + mult * b.y, slope);\n v.z = lact(v.z + mult * b.z, slope);\n v.w = lact(v.w + mult * b.w, slope);\n y4[i] = v;\n }\n}\n\ntorch::Tensor linear_act(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double multiplier, double negative_slope) {\n TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(), \"CUDA tensors required\");\n TORCH_CHECK(x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat && bias.scalar_type() == at::kFloat, \"float32 only\");\n const int64_t M64 = x.size(0), K64 = x.size(1), N64 = weight.size(0);\n auto y = torch::empty({M64, N64}, x.options());\n int M = (int)M64, K = (int)K64, N = (int)N64;\n\n cublasLtHandle_t handle = get_lt_handle();\n cublasLtMatmulDesc_t opDesc;\n cublasLtMatrixLayout_t aDesc, bDesc, cDesc;\n CUBLAS_CHECK(cublasLtMatmulDescCreate(&opDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F));\n cublasOperation_t transa = CUBLAS_OP_T;\n cublasOperation_t transb = CUBLAS_OP_N;\n CUBLAS_CHECK(cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa)));\n CUBLAS_CHECK(cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transb)));\n CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&aDesc, CUDA_R_32F, K, N, K));\n CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&bDesc, CUDA_R_32F, K, M, K));\n CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&cDesc, CUDA_R_32F, N, M, N));\n\n float alpha = (float)multiplier;\n float beta = 0.0f;\n cublasLtMatmulPreference_t pref;\n CUBLAS_CHECK(cublasLtMatmulPreferenceCreate(&pref));\n size_t workspaceSize = 0;\n CUBLAS_CHECK(cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize)));\n cublasLtMatmulHeuristicResult_t heuristic;\n int returned = 0;\n CUBLAS_CHECK(cublasLtMatmulAlgoGetHeuristic(handle, opDesc, aDesc, bDesc, cDesc, cDesc, pref, 1, &heuristic, &returned));\n TORCH_CHECK(returned > 0, \"no cublasLt algo\");\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n CUBLAS_CHECK(cublasLtMatmul(handle, opDesc, &alpha,\n weight.data_ptr<float>(), aDesc,\n x.data_ptr<float>(), bDesc,\n &beta,\n y.data_ptr<float>(), cDesc,\n y.data_ptr<float>(), cDesc,\n &heuristic.algo, nullptr, 0, stream));\n CUBLAS_CHECK(cublasLtMatmulPreferenceDestroy(pref));\n CUBLAS_CHECK(cublasLtMatrixLayoutDestroy(aDesc));\n CUBLAS_CHECK(cublasLtMatrixLayoutDestroy(bDesc));\n CUBLAS_CHECK(cublasLtMatrixLayoutDestroy(cDesc));\n CUBLAS_CHECK(cublasLtMatmulDescDestroy(opDesc));\n\n long long n = y.numel();\n if (n > 0) {\n long long n4 = n >> 2;\n int threads = 256;\n int blocks = (int)((n4 + threads - 1) / threads);\n if (blocks < 1) blocks = 1;\n if (blocks > 65535) blocks = 65535;\n add_bias_lrelu_vec4<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), bias.data_ptr<float>(), n4, N, (float)multiplier, (float)negative_slope);\n }\n return y;\n}\n'''\n_ext = load_inline(name='linear_act_ext_lt1', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['linear_act'], extra_cflags=['-O3'], extra_cuda_cflags=['-O3'], extra_ldflags=['-lcublasLt','-lcublas'], verbose=False)\nclass ModelNew(nn.Module):\n def __init__(self, in_features, out_features, multiplier, negative_slope):\n super(ModelNew, self).__init__()\n self.gemm = nn.Linear(in_features, out_features)\n self.multiplier = float(multiplier)\n self.negative_slope = float(negative_slope)\n def forward(self, x):\n return _ext.linear_act(x, self.gemm.weight, self.gemm.bias, self.multiplier, self.negative_slope)\n"
}