submission 115610

import torch
from torch.utils.cpp_extension import load_inline

# --- 1. CUDA Kernel Source ---
cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <algorithm>

// 辅助函数：float4 累加
template <int OC>
__device__ __forceinline__ void accumulate_oc(const float* wptr, float v, float* acc) {
  if constexpr (OC % 4 == 0) {
#pragma unroll
    for (int oc = 0; oc < OC; oc += 4) {
      const float4 w = reinterpret_cast<const float4*>(wptr)[oc / 4];
      acc[oc + 0] = fmaf(v, w.x, acc[oc + 0]);
      acc[oc + 1] = fmaf(v, w.y, acc[oc + 1]);
      acc[oc + 2] = fmaf(v, w.z, acc[oc + 2]);
      acc[oc + 3] = fmaf(v, w.w, acc[oc + 3]);
    }
  } else {
#pragma unroll
    for (int oc = 0; oc < OC; ++oc) {
      acc[oc] = fmaf(v, wptr[oc], acc[oc]);
    }
  }
}

namespace {

template <int TW, int TH>
inline dim3 make_grid(int out_w, int out_h, int batch, int channel_chunk) {
  const int grid_x = (out_w + TW - 1) / TW;
  const int grid_y = (out_h + TH - 1) / TH;
  const int grid_z = batch * channel_chunk;
  return dim3(grid_x, grid_y, grid_z);
}

// ---- 专门化多通道静态共享内存核 (保持不变，这部分逻辑是好的) ----
template <int K, int C, int OC, int TW, int TH>
__launch_bounds__(TW * TH, 2)
__global__ void conv2d_multi_static(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    float* __restrict__ output,
    int batch,
    int out_channels,
    int height,
    int width) {
  
  // ... 保持你的原始逻辑 ...
  // 为了节省篇幅，这里省略重复代码，但编译时需要包含你的原始 conv2d_multi_static 实现
  // 请确保这里包含你原代码中该函数的完整内容
  
  static_assert(TW * TH <= 1024, "block too large");
  constexpr int PATCH_W = TW + K - 1;
  constexpr int PATCH_H = TH + K - 1;
  constexpr int PATCH_STRIDE = (PATCH_W >= 15) ? PATCH_W : (PATCH_W + 1);

  const int out_w = width - K + 1;
  const int out_h = height - K + 1;

  const int tile_x = blockIdx.x * TW;
  const int tile_y = blockIdx.y * TH;
  if (tile_x >= out_w || tile_y >= out_h) return;

  const int channel_chunks = (out_channels + OC - 1) / OC;
  const int bc = static_cast<int>(blockIdx.z);
  const int b_idx = bc / channel_chunks;
  const int chunk_idx = bc - b_idx * channel_chunks;
  const int c_start = chunk_idx * OC;
  if (b_idx >= batch || c_start >= out_channels) return;

  alignas(16) __shared__ float k_smem[C][K][K][OC];
  alignas(16) __shared__ float in_smem[C][PATCH_H][PATCH_STRIDE];

  const int tid = threadIdx.y * blockDim.x + threadIdx.x;
  const int blk_threads = blockDim.x * blockDim.y;

  const int total_k = OC * C * K * K;
  for (int idx = tid; idx < total_k; idx += blk_threads) {
    const int tmp = idx / OC;
    const int oc = idx - tmp * OC;
    const int ci = tmp / (K * K);
    const int rem0 = tmp - ci * K * K;
    const int kh = rem0 / K;
    const int kw = rem0 - kh * K;
    const int g_oc = c_start + oc;
    if (g_oc < out_channels) { 
        const int k_idx = ((g_oc * C + ci) * K + kh) * K + kw;
        k_smem[ci][kh][kw][oc] = weight[k_idx];
    } else {
        k_smem[ci][kh][kw][oc] = 0.0f;
    }
  }

  const int patch_elems = C * PATCH_H * PATCH_W;
  for (int idx = tid; idx < patch_elems; idx += blk_threads) {
    const int ci = idx / (PATCH_H * PATCH_W);
    const int rem0 = idx - ci * PATCH_H * PATCH_W;
    const int py = rem0 / PATCH_W;
    const int px = rem0 - py * PATCH_W;
    const int gx = tile_x + px;
    const int gy = tile_y + py;
    float val = 0.f;
    if (gx < width && gy < height && b_idx < batch) {
      const int in_idx = ((b_idx * C + ci) * height + gy) * width + gx;
      val = input[in_idx];
    }
    in_smem[ci][py][px] = val;
  }
  __syncthreads();

  const int out_x = tile_x + threadIdx.x;
  const int out_y = tile_y + threadIdx.y;
  if (out_x < out_w && out_y < out_h) {
    float acc[OC];
#pragma unroll
    for (int oc = 0; oc < OC; ++oc) acc[oc] = 0.f;

#pragma unroll
    for (int ci = 0; ci < C; ++ci) {
      const float* in_ptr = &in_smem[ci][threadIdx.y][threadIdx.x];
#pragma unroll
      for (int kh = 0; kh < K; ++kh) {
#pragma unroll
        for (int kw = 0; kw < K; ++kw) {
          const float v = in_ptr[kh * PATCH_STRIDE + kw];
          const float* wptr = &k_smem[ci][kh][kw][0];
          accumulate_oc<OC>(wptr, v, acc);
        }
      }
    }

#pragma unroll
    for (int oc = 0; oc < OC; ++oc) {
      const int g_oc = c_start + oc;
      if (g_oc < out_channels) {
        const int out_idx = ((b_idx * out_channels + g_oc) * out_h + out_y) * out_w + out_x;
        output[out_idx] = acc[oc];
      }
    }
  }
}

// ---- 【修复版】通用 fallback：Global Memory 版本 ----
// 不使用 Shared Memory，避免因 channel 过多导致 Invalid Argument
__global__ void conv2d_generic_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    float* __restrict__ output,
    int batch,
    int in_channels,
    int out_channels,
    int height,
    int width,
    int kernel_size) {
  
  // 每个 block 处理 16x16 的输出区域
  const int out_w = width - kernel_size + 1;
  const int out_h = height - kernel_size + 1;
  
  const int col = blockIdx.x * blockDim.x + threadIdx.x;
  const int row = blockIdx.y * blockDim.y + threadIdx.y;
  
  // Z 轴映射到 Batch 和 Out Channel
  const int bc = blockIdx.z;
  const int b = bc / out_channels;
  const int c_out = bc % out_channels; // 修正逻辑：bc = b * out_c + c_out

  if (col >= out_w || row >= out_h || b >= batch) return;

  float acc = 0.0f;

  // 直接遍历输入通道，利用 L1 Cache 进行缓存
  for (int c = 0; c < in_channels; ++c) {
    for (int kh = 0; kh < kernel_size; ++kh) {
      int in_h = row + kh;
      for (int kw = 0; kw < kernel_size; ++kw) {
        int in_w = col + kw;
        
        // 计算全局索引
        int in_idx = ((b * in_channels + c) * height + in_h) * width + in_w;
        int k_idx = ((c_out * in_channels + c) * kernel_size + kh) * kernel_size + kw;
        
        // 乘累加
        acc += input[in_idx] * weight[k_idx];
      }
    }
  }

  int out_idx = ((b * out_channels + c_out) * out_h + row) * out_w + col;
  output[out_idx] = acc;
}

}  // namespace

// --- C++ Launcher ---
void conv2d_forward_wrapper(
    torch::Tensor input,
    torch::Tensor kernel,
    torch::Tensor output) {
  const int batch = static_cast<int>(input.size(0));
  const int in_channels = static_cast<int>(input.size(1));
  const int height = static_cast<int>(input.size(2));
  const int width = static_cast<int>(input.size(3));

  const int out_channels = static_cast<int>(kernel.size(0));
  const int kernel_size = static_cast<int>(kernel.size(2));

  const int out_h = height - kernel_size + 1;
  const int out_w = width - kernel_size + 1;

  // --- 专门化路径 (Fast Path) ---
  // K4 C8 OC8
  if (kernel_size == 4 && in_channels == 8 && out_channels == 8) {
    const int channel_chunks = (out_channels + 8 - 1) / 8;
    const dim3 grid = make_grid<12, 8>(out_w, out_h, batch, channel_chunks);
    const dim3 block(12, 8);
    conv2d_multi_static<4, 8, 8, 12, 8><<<grid, block>>>(
        input.data_ptr<float>(), kernel.data_ptr<float>(), output.data_ptr<float>(),
        batch, out_channels, height, width);
    return;
  }
  // K6 C16 OC16
  if (kernel_size == 6 && in_channels == 16 && out_channels == 16) {
    const int channel_chunks = (out_channels + 4 - 1) / 4;
    const dim3 grid = make_grid<16, 10>(out_w, out_h, batch, channel_chunks);
    const dim3 block(16, 10);
    conv2d_multi_static<6, 16, 4, 16, 10><<<grid, block>>>(
        input.data_ptr<float>(), kernel.data_ptr<float>(), output.data_ptr<float>(),
        batch, out_channels, height, width);
    return;
  }
  // K8 C8 OC8
  if (kernel_size == 8 && in_channels == 8 && out_channels == 8) {
    const int channel_chunks = (out_channels + 4 - 1) / 4;
    const dim3 grid = make_grid<16, 10>(out_w, out_h, batch, channel_chunks);
    const dim3 block(16, 10);
    conv2d_multi_static<8, 8, 4, 16, 10><<<grid, block>>>(
        input.data_ptr<float>(), kernel.data_ptr<float>(), output.data_ptr<float>(),
        batch, out_channels, height, width);
    return;
  }
  // K8 C16 OC16
  if (kernel_size == 8 && in_channels == 16 && out_channels == 16) {
    const int channel_chunks = (out_channels + 2 - 1) / 2;
    const dim3 grid = make_grid<16, 8>(out_w, out_h, batch, channel_chunks);
    const dim3 block(16, 8);
    conv2d_multi_static<8, 16, 2, 16, 8><<<grid, block>>>(
        input.data_ptr<float>(), kernel.data_ptr<float>(), output.data_ptr<float>(),
        batch, out_channels, height, width);
    return;
  }

  // --- 通用 Fallback (Safe Path) ---
  // 使用 Global Memory 版本，不申请 Shared Memory，避免 crash
  dim3 block(16, 16);
  dim3 grid(
      (out_w + block.x - 1) / block.x,
      (out_h + block.y - 1) / block.y,
      batch * out_channels // Z 维度处理 Batch 和 OC
  );
  
  // 移除 smem_bytes 参数
  conv2d_generic_kernel<<<grid, block>>>(
      input.data_ptr<float>(),
      kernel.data_ptr<float>(),
      output.data_ptr<float>(),
      batch,
      in_channels,
      out_channels,
      height,
      width,
      kernel_size);
}
"""

cpp_source = (
    "void conv2d_forward_wrapper("
    "    torch::Tensor input,"
    "    torch::Tensor kernel,"
    "    torch::Tensor output"
    ");"
)

conv2d_ext = load_inline(
    name="conv2d_v5_robust",
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=["conv2d_forward_wrapper"],
    extra_cflags=["-std=c++17", "-O3"],
    extra_cuda_cflags=[
        "-O3", "--use_fast_math", "-std=c++17", 
        "-U__CUDA_NO_HALF_OPERATORS__", "-U__CUDA_NO_HALF_CONVERSIONS__"
    ],
    verbose=False,
)

def custom_kernel(input_tuple):
    input_tensor, kernel, output = input_tuple
    if not input_tensor.is_contiguous(): input_tensor = input_tensor.contiguous()
    if not kernel.is_contiguous(): kernel = kernel.contiguous()
    if not output.is_contiguous(): output = output.contiguous()
    
    conv2d_ext.conv2d_forward_wrapper(input_tensor, kernel, output)
    return output