callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_src = r'''\n#include <torch/extension.h>\ntorch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias,\n torch::Tensor post_bias, double scale_d,\n int64_t stride_l, int64_t padding_l, int64_t output_padding_l);\n'''\n\ncuda_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <vector>\n#include <cfloat>\n\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be CUDA\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::kFloat, #x \" must be float\")\n\n__global__ void pack_w_kernel(const float* __restrict__ w, float* __restrict__ wp,\n int IC, int OC, int K) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = IC * K * K * OC;\n if (idx >= total) return;\n int oc = idx % OC;\n int t = idx / OC;\n int kw = t % K;\n t /= K;\n int kh = t % K;\n int ic = t / K;\n wp[idx] = w[((ic * OC + oc) * K + kh) * K + kw];\n}\n\n__global__ __launch_bounds__(128, 2)\nvoid fused_deconv_softmax_kernel(const float* __restrict__ x,\n const float* __restrict__ wp,\n const float* __restrict__ conv_b,\n const float* __restrict__ post_b,\n float* __restrict__ out,\n int N, int IC, int H, int W, int OC, int K,\n int stride, int pad, int Ho, int Wo,\n float scale) {\n int oc = threadIdx.x;\n int sidx = blockIdx.x;\n int ox = sidx % Wo;\n int oy = (sidx / Wo) % Ho;\n int n = sidx / (Wo * Ho);\n float acc = conv_b[oc];\n for (int ic = 0; ic < IC; ++ic) {\n #pragma unroll\n for (int kh = 0; kh < 4; ++kh) {\n int iy_num = oy + pad - kh;\n if (iy_num < 0 || iy_num % stride != 0) continue;\n int iy = iy_num / stride;\n if (iy < 0 || iy >= H) continue;\n #pragma unroll\n for (int kw = 0; kw < 4; ++kw) {\n int ix_num = ox + pad - kw;\n if (ix_num < 0 || ix_num % stride != 0) continue;\n int ix = ix_num / stride;\n if (ix < 0 || ix >= W) continue;\n float xv = x[((n * IC + ic) * H + iy) * W + ix];\n float wv = wp[(((ic * K + kh) * K + kw) * OC) + oc];\n acc = fmaf(xv, wv, acc);\n }\n }\n }\n __shared__ float smem[128];\n smem[oc] = acc;\n __syncthreads();\n for (int offset = 64; offset > 0; offset >>= 1) {\n if (oc < offset) smem[oc] = fmaxf(smem[oc], smem[oc + offset]);\n __syncthreads();\n }\n float m = smem[0];\n float e = __expf(acc - m);\n smem[oc] = e;\n __syncthreads();\n for (int offset = 64; offset > 0; offset >>= 1) {\n if (oc < offset) smem[oc] += smem[oc + offset];\n __syncthreads();\n }\n float soft = e / smem[0];\n float z = (soft + post_b[oc]) * scale;\n out[((n * OC + oc) * Ho + oy) * Wo + ox] = 1.0f / (1.0f + __expf(-z));\n}\n\ntorch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias,\n torch::Tensor post_bias, double scale_d,\n int64_t stride_l, int64_t padding_l, int64_t output_padding_l) {\n CHECK_CUDA(x); CHECK_CUDA(weight); CHECK_CUDA(conv_bias); CHECK_CUDA(post_bias);\n CHECK_CONTIGUOUS(x); CHECK_CONTIGUOUS(weight); CHECK_CONTIGUOUS(conv_bias); CHECK_CONTIGUOUS(post_bias);\n CHECK_FLOAT(x); CHECK_FLOAT(weight); CHECK_FLOAT(conv_bias); CHECK_FLOAT(post_bias);\n int N = x.size(0), IC = x.size(1), H = x.size(2), W = x.size(3);\n int OC = weight.size(1), K = weight.size(2);\n int stride = (int)stride_l, pad = (int)padding_l, outpad = (int)output_padding_l;\n int Ho = (H - 1) * stride - 2 * pad + K + outpad;\n int Wo = (W - 1) * stride - 2 * pad + K + outpad;\n auto out = torch::empty({N, OC, Ho, Wo}, x.options());\n auto wp = torch::empty({IC, K, K, OC}, x.options());\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n int totalw = IC * K * K * OC;\n pack_w_kernel<<<(totalw + 255)/256, 256, 0, stream>>>(weight.data_ptr<float>(), wp.data_ptr<float>(), IC, OC, K);\n int blocks = N * Ho * Wo;\n fused_deconv_softmax_kernel<<<blocks, 128, 0, stream>>>(x.data_ptr<float>(), wp.data_ptr<float>(), conv_bias.data_ptr<float>(), post_bias.data_ptr<float>(), out.data_ptr<float>(), N, IC, H, W, OC, K, stride, pad, Ho, Wo, (float)scale_d);\n return out;\n}\n'''\n\n_ext = load_inline(name='fused_deconv_softmax_ext_v2', cpp_sources=cpp_src, cuda_sources=cuda_src,\n functions=['fused_forward'], 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, bias_shape, scaling_factor):\n super(ModelNew, self).__init__()\n self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)\n self.bias = nn.Parameter(torch.randn(bias_shape))\n self.scaling_factor = scaling_factor\n self.stride = stride if isinstance(stride, int) else stride[0]\n self.padding = padding if isinstance(padding, int) else padding[0]\n self.output_padding = output_padding if isinstance(output_padding, int) else output_padding[0]\n def forward(self, x):\n return _ext.fused_forward(x.contiguous(), self.conv_transpose.weight.contiguous(), self.conv_transpose.bias.contiguous(), self.bias.contiguous(), float(self.scaling_factor), int(self.stride), int(self.padding), int(self.output_padding))\n"
}