callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n_cpp = r'''\n#include <torch/extension.h>\ntorch::Tensor causal_conv1d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t dilation, bool has_bias);\n'''\n_cuda = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#define BM 32\n#define BN 32\n#define BK 32\n__device__ __forceinline__ float load_A_val(const float* __restrict__ x, int m, int kk, int B, int Cin, int L, int K, int dilation, int pad) {\n int b = m / L;\n int t = m - b * L;\n int ci = kk / K;\n int kpos = kk - ci * K;\n int it = t + kpos * dilation - pad;\n if ((unsigned)b < (unsigned)B && (unsigned)it < (unsigned)L) return x[((b * Cin + ci) * L) + it];\n return 0.0f;\n}\n__global__ void causal_conv1d_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ bias, float* __restrict__ y, int B, int Cin, int Cout, int L, int K, int dilation, int has_bias) {\n __shared__ float As[BM][BK + 1];\n __shared__ float Bs[BK][BN + 1];\n int tx=threadIdx.x, ty=threadIdx.y, tid=ty*16+tx;\n int m0=blockIdx.y*BM+ty, m1=m0+16, n0=blockIdx.x*BN+tx, n1=n0+16;\n int M=B*L, Ktot=Cin*K, pad=(K-1)*dilation;\n float acc00=0.f,acc01=0.f,acc10=0.f,acc11=0.f;\n for(int kk0=0; kk0<Ktot; kk0+=BK){\n#pragma unroll\n for(int i=0;i<4;++i){\n int idx=tid+i*256; int r=idx>>5; int c=idx&31; int m=blockIdx.y*BM+r; int kk=kk0+c;\n float av=0.f; if(m<M && kk<Ktot) av=load_A_val(x,m,kk,B,Cin,L,K,dilation,pad); As[r][c]=av;\n int n=blockIdx.x*BN+c; int kglob=kk0+r; float bv=0.f; if(n<Cout && kglob<Ktot){ int ci=kglob/K; int kpos=kglob-ci*K; bv=w[(n*Cin+ci)*K+kpos]; } Bs[r][c]=bv;\n }\n __syncthreads();\n#pragma unroll\n for(int p=0;p<BK;++p){ float a0=As[ty][p], a1=As[ty+16][p], b0=Bs[p][tx], b1=Bs[p][tx+16]; acc00 += a0*b0; acc01 += a0*b1; acc10 += a1*b0; acc11 += a1*b1; }\n __syncthreads();\n }\n if(m0<M){int b=m0/L,t=m0-b*L; if(n0<Cout)y[(b*Cout+n0)*L+t]=acc00+(has_bias?bias[n0]:0.f); if(n1<Cout)y[(b*Cout+n1)*L+t]=acc01+(has_bias?bias[n1]:0.f);} \n if(m1<M){int b=m1/L,t=m1-b*L; if(n0<Cout)y[(b*Cout+n0)*L+t]=acc10+(has_bias?bias[n0]:0.f); if(n1<Cout)y[(b*Cout+n1)*L+t]=acc11+(has_bias?bias[n1]:0.f);} \n}\ntorch::Tensor causal_conv1d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t dilation64, bool has_bias){\n x=x.contiguous(); weight=weight.contiguous(); int B=x.size(0),Cin=x.size(1),L=x.size(2),Cout=weight.size(0),K=weight.size(2),dilation=(int)dilation64; auto y=torch::empty({B,Cout,L},x.options()); dim3 block(16,16), grid((Cout+BN-1)/BN,(B*L+BM-1)/BM); causal_conv1d_kernel<<<grid,block,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),weight.data_ptr<float>(),has_bias?bias.data_ptr<float>():nullptr,y.data_ptr<float>(),B,Cin,Cout,L,K,dilation,has_bias?1:0); return y; }\n'''\n_ext=load_inline(name='causal_conv1d_ext_v2',cpp_sources=_cpp,cuda_sources=_cuda,functions=['causal_conv1d_forward'],verbose=False,extra_cuda_cflags=['-O3'],extra_cflags=['-O3'])\nclass ModelNew(nn.Module):\n def __init__(self,in_channels:int,out_channels:int,kernel_size:int,dilation:int=1,bias:bool=False):\n super().__init__(); self.causal_padding=(kernel_size-1)*dilation; self.dilation=dilation; self.conv1d=nn.Conv1d(in_channels,out_channels,kernel_size,dilation=dilation,bias=bias)\n def forward(self,x:torch.Tensor)->torch.Tensor:\n b=self.conv1d.bias\n if b is None:\n b=torch.empty(0,device=x.device,dtype=x.dtype); return _ext.causal_conv1d_forward(x,self.conv1d.weight,b,self.dilation,False)\n return _ext.causal_conv1d_forward(x,self.conv1d.weight,b,self.dilation,True)\n"
}