callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_src = \"\"\"\n#include <torch/extension.h>\ntorch::Tensor convT2d_forward(torch::Tensor x, torch::Tensor weight);\n\"\"\"\n\ncuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n\n__global__ void convT2d_group16_kernel(const float* __restrict__ x,\n const float* __restrict__ w,\n float* __restrict__ y) {\n const int ow = blockIdx.x * 32 + threadIdx.x;\n const int row = blockIdx.y * 4 + threadIdx.y; // N*groups*odd_h\n if (ow >= 766 || row >= 16*4*128) return;\n\n const int odd_h_idx = row & 127;\n int tmp = row >> 7;\n const int g = tmp & 3;\n const int n = tmp >> 2;\n const int oh = (odd_h_idx << 1) + 1;\n const int base_ih = odd_h_idx + 1;\n const int ic0 = g << 3;\n const int oc_base = g << 4;\n\n const int t = ow + 2;\n const int kw0 = t % 3;\n\n float acc0=0.f, acc1=0.f, acc2=0.f, acc3=0.f;\n float acc4=0.f, acc5=0.f, acc6=0.f, acc7=0.f;\n float acc8=0.f, acc9=0.f, acc10=0.f, acc11=0.f;\n float acc12=0.f, acc13=0.f, acc14=0.f, acc15=0.f;\n\n #pragma unroll\n for (int kh = 0; kh < 3; ++kh) {\n const int ih = base_ih - kh;\n if ((unsigned)ih < 128u) {\n #pragma unroll\n for (int kk = 0; kk < 2; ++kk) {\n const int kw = kw0 + kk * 3;\n if (kw < 5) {\n const int iw = (t - kw) / 3;\n if ((unsigned)iw < 256u) {\n const float* xp = x + ((n * 32 + ic0) * 128 + ih) * 256 + iw;\n const float* wp = w + (ic0 * 16 * 15 + kh * 5 + kw);\n #pragma unroll\n for (int ci = 0; ci < 8; ++ci) {\n float xv = xp[ci * 32768];\n const float* wc = wp + ci * 240;\n acc0 = fmaf(xv, wc[0*15], acc0);\n acc1 = fmaf(xv, wc[1*15], acc1);\n acc2 = fmaf(xv, wc[2*15], acc2);\n acc3 = fmaf(xv, wc[3*15], acc3);\n acc4 = fmaf(xv, wc[4*15], acc4);\n acc5 = fmaf(xv, wc[5*15], acc5);\n acc6 = fmaf(xv, wc[6*15], acc6);\n acc7 = fmaf(xv, wc[7*15], acc7);\n acc8 = fmaf(xv, wc[8*15], acc8);\n acc9 = fmaf(xv, wc[9*15], acc9);\n acc10 = fmaf(xv, wc[10*15], acc10);\n acc11 = fmaf(xv, wc[11*15], acc11);\n acc12 = fmaf(xv, wc[12*15], acc12);\n acc13 = fmaf(xv, wc[13*15], acc13);\n acc14 = fmaf(xv, wc[14*15], acc14);\n acc15 = fmaf(xv, wc[15*15], acc15);\n }\n }\n }\n }\n }\n }\n float* yp = y + ((n * 64 + oc_base) * 257 + oh) * 766 + ow;\n const int cstride = 257*766;\n yp[0*cstride]=acc0; yp[1*cstride]=acc1; yp[2*cstride]=acc2; yp[3*cstride]=acc3;\n yp[4*cstride]=acc4; yp[5*cstride]=acc5; yp[6*cstride]=acc6; yp[7*cstride]=acc7;\n yp[8*cstride]=acc8; yp[9*cstride]=acc9; yp[10*cstride]=acc10; yp[11*cstride]=acc11;\n yp[12*cstride]=acc12; yp[13*cstride]=acc13; yp[14*cstride]=acc14; yp[15*cstride]=acc15;\n}\n\ntorch::Tensor convT2d_forward(torch::Tensor x, torch::Tensor weight) {\n CHECK_CUDA(x); CHECK_CUDA(weight);\n CHECK_CONTIGUOUS(x); CHECK_CONTIGUOUS(weight);\n auto y = torch::empty({x.size(0), 64, 257, 766}, x.options());\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n cudaMemsetAsync(y.data_ptr<float>(), 0, y.numel() * sizeof(float), stream);\n dim3 block(32, 4);\n dim3 grid(24, (int)((x.size(0) * 4 * 128 + 3) / 4));\n convT2d_group16_kernel<<<grid, block, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), y.data_ptr<float>());\n C10_CUDA_KERNEL_LAUNCH_CHECK();\n return y;\n}\n'''\n\n_ext = load_inline(name='convT2d_special_ext_v3', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['convT2d_forward'], verbose=False, extra_cuda_cflags=['-O3', '--use_fast_math'])\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple, stride: tuple = (1, 1), padding: tuple = (0, 0), dilation: tuple = (1, 1), groups: int = 1, bias: bool = False):\n super(ModelNew, self).__init__()\n self.conv_transpose2d = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return _ext.convT2d_forward(x.contiguous(), self.conv_transpose2d.weight.contiguous())\n"
}