callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_source = r'''\n#include <torch/extension.h>\ntorch::Tensor fused_convT_softmax_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor scale, double clamp_min, double clamp_max);\n'''\n\ncuda_source = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <float.h>\n\n__global__ void convT3d_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ bias,\n float* __restrict__ y, long long total, int N, int IC, int ID, int IH, int IW, int OC) {\n long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= total) return;\n int OW = IW * 2; int OH = IH * 2; int OD = ID * 2;\n int ow = idx % OW;\n long long t = idx / OW;\n int oh = t % OH; t /= OH;\n int od = t % OD; t /= OD;\n int oc = t % OC; int n = t / OC;\n float acc = bias[oc];\n // stride=2, padding=1, kernel=3. valid k have same parity as out+1\n int kd0 = (od + 1) & 1; // 0 for odd od, 1 for even od\n int kh0 = (oh + 1) & 1;\n int kw0 = (ow + 1) & 1;\n #pragma unroll\n for (int ic = 0; ic < 32; ++ic) {\n #pragma unroll\n for (int kd = kd0; kd < 3; kd += 2) {\n int id = (od + 1 - kd) >> 1;\n if ((unsigned)id >= (unsigned)ID) continue;\n #pragma unroll\n for (int kh = kh0; kh < 3; kh += 2) {\n int ih = (oh + 1 - kh) >> 1;\n if ((unsigned)ih >= (unsigned)IH) continue;\n #pragma unroll\n for (int kw = kw0; kw < 3; kw += 2) {\n int iw = (ow + 1 - kw) >> 1;\n if ((unsigned)iw >= (unsigned)IW) continue;\n float xv = x[(((long long)n*IC + ic)*ID + id)*IH*IW + ih*IW + iw];\n float ww = w[(((ic*OC + oc)*3 + kd)*3 + kh)*3 + kw];\n acc = fmaf(xv, ww, acc);\n }\n }\n }\n }\n y[idx] = acc;\n}\n\ntemplate<int BLOCK>\n__global__ void softmax_clamp_scale_kernel(float* __restrict__ x, const float* __restrict__ scale,\n int rows, int cols, float clamp_min, float clamp_max) {\n int row = blockIdx.x;\n int tid = threadIdx.x;\n int c = row & 63;\n float* base = x + (long long)row * cols;\n float local_max = -FLT_MAX;\n for (int i = tid; i < cols; i += BLOCK) {\n float v = base[i]; v = fminf(fmaxf(v, clamp_min), clamp_max); local_max = fmaxf(local_max, v);\n }\n __shared__ float smem[BLOCK]; smem[tid] = local_max; __syncthreads();\n for (int s=BLOCK/2; s>0; s>>=1) { if (tid<s) smem[tid]=fmaxf(smem[tid],smem[tid+s]); __syncthreads(); }\n float m=smem[0]; float local_sum=0.f;\n for (int i = tid; i < cols; i += BLOCK) { float v=base[i]; v=fminf(fmaxf(v,clamp_min),clamp_max); float e=__expf(v-m); base[i]=e; local_sum += e; }\n smem[tid]=local_sum; __syncthreads();\n for (int s=BLOCK/2; s>0; s>>=1) { if (tid<s) smem[tid]+=smem[tid+s]; __syncthreads(); }\n float mul=scale[c]/smem[0];\n for (int i=tid; i<cols; i+=BLOCK) base[i]*=mul;\n}\n\ntorch::Tensor fused_convT_softmax_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor scale, double clamp_min, double clamp_max) {\n int N=x.size(0), IC=x.size(1), ID=x.size(2), IH=x.size(3), IW=x.size(4); int OC=weight.size(1);\n auto y = torch::empty({N, OC, ID*2, IH*2, IW*2}, x.options());\n long long total = y.numel();\n int threads=256; int blocks=(total+threads-1)/threads;\n auto stream=at::cuda::getCurrentCUDAStream();\n convT3d_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), y.data_ptr<float>(), total, N, IC, ID, IH, IW, OC);\n int rows=N*OC; int cols=ID*2*IH*2*IW*2;\n softmax_clamp_scale_kernel<256><<<rows, 256, 0, stream>>>(y.data_ptr<float>(), scale.data_ptr<float>(), rows, cols, (float)clamp_min, (float)clamp_max);\n return y;\n}\n'''\n_ext = load_inline(name='fused_convT_softmax_ext_v1', cpp_sources=cpp_source, cuda_sources=cuda_source, functions=['fused_convT_softmax_cuda'], verbose=False, extra_cuda_cflags=['-O3', '--use_fast_math'])\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, pool_kernel_size, clamp_min, clamp_max):\n super(ModelNew, self).__init__()\n self.avg_pool = nn.AvgPool3d(pool_kernel_size)\n self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)\n self.clamp_min=clamp_min; self.clamp_max=clamp_max\n self.scale=nn.Parameter(torch.ones(1,out_channels,1,1,1))\n def forward(self,x):\n x=self.avg_pool(x)\n return _ext.fused_convT_softmax_cuda(x, self.conv_transpose.weight, self.conv_transpose.bias, self.scale, self.clamp_min, self.clamp_max)\n"
}