callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_sources = r'''\n#include <torch/extension.h>\ntorch::Tensor fused_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor add_bias);\n'''\n\ncuda_sources = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda_runtime.h>\n#include <cmath>\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#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::kFloat, #x \" must be float32\")\n\n__inline__ __device__ float warp_sum(float v) {\n for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off);\n return v;\n}\n__inline__ __device__ float warp_max(float v) {\n for (int off = 16; off > 0; off >>= 1) v = fmaxf(v, __shfl_down_sync(0xffffffff, v, off));\n return v;\n}\n\n__global__ void reduce_x_kernel(const float* __restrict__ x, float* __restrict__ xs,\n int B, int Cin, int H, int W) {\n int bc = blockIdx.x;\n int c = bc % Cin;\n int b = bc / Cin;\n int HW = H * W;\n const float* base = x + ((long long)b * Cin + c) * (long long)HW;\n float sum = 0.0f;\n int tid = threadIdx.x;\n if ((HW & 3) == 0) {\n const float4* p4 = reinterpret_cast<const float4*>(base);\n int N4 = HW >> 2;\n for (int i = tid; i < N4; i += 256) {\n float4 v = p4[i];\n sum += (v.x + v.y) + (v.z + v.w);\n }\n } else {\n for (int i = tid; i < HW; i += 256) sum += base[i];\n }\n sum = warp_sum(sum);\n __shared__ float warp_part[8];\n int lane = tid & 31;\n int warp = tid >> 5;\n if (lane == 0) warp_part[warp] = sum;\n __syncthreads();\n float block_sum = (tid < 8) ? warp_part[lane] : 0.0f;\n if (warp == 0) block_sum = warp_sum(block_sum);\n if (tid == 0) xs[bc] = block_sum;\n}\n\n__global__ void final_lse_64_128_3_kernel(const float* __restrict__ xs,\n const float* __restrict__ weight,\n const float* __restrict__ conv_bias,\n const float* __restrict__ add_bias,\n float* __restrict__ out,\n int H, int W) {\n int b = blockIdx.x;\n int oc = threadIdx.x; // 0..127\n float inv_area = 1.0f / ((float)(H + 2) * (float)(W + 2));\n float dot = 0.0f;\n#pragma unroll\n for (int ic = 0; ic < 64; ++ic) {\n const float* wb = weight + (((ic * 128 + oc) * 3) * 3);\n float ws = ((wb[0] + wb[1]) + (wb[2] + wb[3])) + ((wb[4] + wb[5]) + (wb[6] + wb[7] + wb[8]));\n dot += xs[b * 64 + ic] * ws;\n }\n float v = conv_bias[oc] + add_bias[oc] + dot * inv_area;\n __shared__ float vals[128];\n __shared__ float warp_part[4];\n vals[oc] = v;\n float m = warp_max(v);\n int lane = oc & 31;\n int warp = oc >> 5;\n if (lane == 0) warp_part[warp] = m;\n __syncthreads();\n float mm = (oc < 4) ? warp_part[lane] : -3.402823466e38f;\n if (warp == 0) mm = warp_max(mm);\n if (oc == 0) warp_part[0] = mm;\n __syncthreads();\n float maxv = warp_part[0];\n float s = expf(vals[oc] - maxv);\n s = warp_sum(s);\n if (lane == 0) warp_part[warp] = s;\n __syncthreads();\n float ss = (oc < 4) ? warp_part[lane] : 0.0f;\n if (warp == 0) ss = warp_sum(ss);\n if (oc == 0) out[b] = 10.0f * (logf(ss) + maxv);\n}\n\n__device__ __forceinline__ float channel_value_generic(int b, int oc, const float* __restrict__ xs,\n const float* __restrict__ weight,\n const float* __restrict__ conv_bias,\n const float* __restrict__ add_bias,\n int Cin, int Cout, int K, float inv_area) {\n float acc = conv_bias[oc] + add_bias[oc];\n float dot = 0.0f;\n for (int ic = 0; ic < Cin; ++ic) {\n const float* wbase = weight + (((long long)ic * Cout + oc) * K * K);\n float ws = 0.0f;\n for (int kk = 0; kk < K * K; ++kk) ws += wbase[kk];\n dot += xs[b * Cin + ic] * ws;\n }\n acc += dot * inv_area;\n return acc;\n}\n\n__global__ void final_lse_kernel(const float* __restrict__ xs,\n const float* __restrict__ weight,\n const float* __restrict__ conv_bias,\n const float* __restrict__ add_bias,\n float* __restrict__ out,\n int B, int Cin, int Cout, int K, int outH, int outW) {\n int b = blockIdx.x;\n int tid = threadIdx.x;\n float inv_area = 1.0f / ((float)outH * (float)outW);\n float local_max = -3.402823466e38f;\n for (int oc = tid; oc < Cout; oc += blockDim.x) {\n float v = channel_value_generic(b, oc, xs, weight, conv_bias, add_bias, Cin, Cout, K, inv_area);\n local_max = fmaxf(local_max, v);\n }\n __shared__ float smem[256];\n smem[tid] = local_max;\n __syncthreads();\n for (int s = blockDim.x >> 1; s > 0; s >>= 1) {\n if (tid < s) smem[tid] = fmaxf(smem[tid], smem[tid + s]);\n __syncthreads();\n }\n float m = smem[0];\n float local_sum = 0.0f;\n for (int oc = tid; oc < Cout; oc += blockDim.x) {\n float v = channel_value_generic(b, oc, xs, weight, conv_bias, add_bias, Cin, Cout, K, inv_area);\n local_sum += expf(v - m);\n }\n smem[tid] = local_sum;\n __syncthreads();\n for (int s = blockDim.x >> 1; s > 0; s >>= 1) {\n if (tid < s) smem[tid] += smem[tid + s];\n __syncthreads();\n }\n if (tid == 0) out[b] = 10.0f * (logf(smem[0]) + m);\n}\n\ntorch::Tensor fused_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor add_bias) {\n CHECK_CUDA(x); CHECK_CUDA(weight); CHECK_CUDA(conv_bias); CHECK_CUDA(add_bias);\n CHECK_CONTIGUOUS(x); CHECK_CONTIGUOUS(weight); CHECK_CONTIGUOUS(conv_bias); CHECK_CONTIGUOUS(add_bias);\n CHECK_FLOAT(x); CHECK_FLOAT(weight); CHECK_FLOAT(conv_bias); CHECK_FLOAT(add_bias);\n int B = (int)x.size(0);\n int Cin = (int)x.size(1);\n int H = (int)x.size(2);\n int W = (int)x.size(3);\n int Cout = (int)weight.size(1);\n int K = (int)weight.size(2);\n int outH = H + K - 1;\n int outW = W + K - 1;\n auto opts = x.options();\n torch::Tensor xs = torch::empty({B, Cin}, opts);\n torch::Tensor out = torch::empty({B, 1}, opts);\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n reduce_x_kernel<<<B * Cin, 256, 0, stream>>>(x.data_ptr<float>(), xs.data_ptr<float>(), B, Cin, H, W);\n if (Cin == 64 && Cout == 128 && K == 3) {\n final_lse_64_128_3_kernel<<<B, 128, 0, stream>>>(xs.data_ptr<float>(), weight.data_ptr<float>(), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), out.data_ptr<float>(), H, W);\n } else {\n final_lse_kernel<<<B, 256, 0, stream>>>(xs.data_ptr<float>(), weight.data_ptr<float>(), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), out.data_ptr<float>(), B, Cin, Cout, K, outH, outW);\n }\n return out;\n}\n'''\n\n_ext = load_inline(name='fused_convtranspose_gap_lse_ext_v3', cpp_sources=cpp_sources, cuda_sources=cuda_sources, functions=['fused_forward_cuda'], verbose=False, extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'])\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, bias_shape):\n super(ModelNew, self).__init__()\n self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size)\n self.bias = nn.Parameter(torch.randn(bias_shape))\n\n def forward(self, x):\n return _ext.fused_forward_cuda(x.contiguous(), self.conv_transpose.weight.contiguous(), self.conv_transpose.bias.contiguous(), self.bias.contiguous())\n"
}