callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\ncpp_src = r'''\n#include <torch/extension.h>\ntorch::Tensor fused_linear_min_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor constant);\n'''\ncuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda_runtime.h>\n#include <cublasLt.h>\n__global__ void minsub_kernel(float* y, const float* cptr, int total){int idx=blockIdx.x*blockDim.x+threadIdx.x; float c=cptr[0]; for(int i=idx;i<total;i+=blockDim.x*gridDim.x){float v=y[i]; y[i]=fminf(v,c)-c;}}\ntorch::Tensor fused_linear_min_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor constant){\n int M=x.size(0),K=x.size(1),N=weight.size(0); auto y=torch::empty({M,N},x.options()); auto workspace=torch::empty({1<<24}, x.options().dtype(torch::kUInt8));\n static cublasLtHandle_t ltHandle=nullptr; if(!ltHandle) cublasLtCreate(<Handle);\n cublasLtMatmulDesc_t op; cublasLtMatrixLayout_t ad,bd,cd; cublasLtMatmulPreference_t pref; cublasOperation_t ta=CUBLAS_OP_N,tb=CUBLAS_OP_T;\n cublasLtMatmulDescCreate(&op,CUBLAS_COMPUTE_32F,CUDA_R_32F); cublasLtMatmulDescSetAttribute(op,CUBLASLT_MATMUL_DESC_TRANSA,&ta,sizeof(ta)); cublasLtMatmulDescSetAttribute(op,CUBLASLT_MATMUL_DESC_TRANSB,&tb,sizeof(tb));\n cublasLtEpilogue_t epi=CUBLASLT_EPILOGUE_BIAS; cublasLtMatmulDescSetAttribute(op,CUBLASLT_MATMUL_DESC_EPILOGUE,&epi,sizeof(epi)); const void* bp=bias.data_ptr<float>(); cublasLtMatmulDescSetAttribute(op,CUBLASLT_MATMUL_DESC_BIAS_POINTER,&bp,sizeof(bp));\n cublasLtOrder_t order=CUBLASLT_ORDER_ROW; cublasLtMatrixLayoutCreate(&ad,CUDA_R_32F,M,K,K); cublasLtMatrixLayoutSetAttribute(ad,CUBLASLT_MATRIX_LAYOUT_ORDER,&order,sizeof(order)); cublasLtMatrixLayoutCreate(&bd,CUDA_R_32F,N,K,K); cublasLtMatrixLayoutSetAttribute(bd,CUBLASLT_MATRIX_LAYOUT_ORDER,&order,sizeof(order)); cublasLtMatrixLayoutCreate(&cd,CUDA_R_32F,M,N,N); cublasLtMatrixLayoutSetAttribute(cd,CUBLASLT_MATRIX_LAYOUT_ORDER,&order,sizeof(order));\n cublasLtMatmulPreferenceCreate(&pref); size_t ws=1<<24; cublasLtMatmulPreferenceSetAttribute(pref,CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,&ws,sizeof(ws)); cublasLtMatmulHeuristicResult_t heur; int ret=0; cublasLtMatmulAlgoGetHeuristic(ltHandle,op,ad,bd,cd,cd,pref,1,&heur,&ret);\n float alpha=1.0f,beta=0.0f; cublasStatus_t stat=CUBLAS_STATUS_NOT_SUPPORTED; if(ret>0) stat=cublasLtMatmul(ltHandle,op,&alpha,x.data_ptr<float>(),ad,weight.data_ptr<float>(),bd,&beta,y.data_ptr<float>(),cd,y.data_ptr<float>(),cd,&heur.algo,workspace.data_ptr(),ws,at::cuda::getCurrentCUDAStream());\n cublasLtMatmulPreferenceDestroy(pref); cublasLtMatmulDescDestroy(op); cublasLtMatrixLayoutDestroy(ad); cublasLtMatrixLayoutDestroy(bd); cublasLtMatrixLayoutDestroy(cd); if(stat!=CUBLAS_STATUS_SUCCESS) C10_CUDA_CHECK(cudaErrorUnknown);\n int total=M*N; minsub_kernel<<<(total+255)/256,256,0,at::cuda::getCurrentCUDAStream()>>>(y.data_ptr<float>(),constant.data_ptr<float>(),total); C10_CUDA_KERNEL_LAUNCH_CHECK(); return y;}\n'''\n_ext=load_inline(name='fused_linear_min_ext_lt2',cpp_sources=cpp_src,cuda_sources=cuda_src,functions=['fused_linear_min_cuda'],extra_cuda_cflags=['-O3'],extra_cflags=['-O3'],verbose=False)\nclass ModelNew(nn.Module):\n def __init__(self,in_features,out_features,constant):\n super(ModelNew,self).__init__(); self.linear=nn.Linear(in_features,out_features); self.constant=nn.Parameter(torch.tensor(constant))\n def forward(self,x): return _ext.fused_linear_min_cuda(x,self.linear.weight,self.linear.bias,self.constant)\n"
}