```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

conv_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <mma.h>
#include <c10/util/Optional.h>

// ---------------------------------------------------------------------------
// Custom implicit-GEMM tensor-core convolution for the exact target config:
//   stride=1, pad=0, dilation=1, groups=1, C=32, OC=64, KH=5, KW=9
// One thread block computes all 64 output channels for 64 output pixels that
// share the same output row (one batch, one oh, 64-wide ow tile).
// K dimension (C*KH*KW = 1440) is processed per input channel: weights tile
// (64x48) and input patch tile (48x64) are staged in shared memory and fed to
// WMMA (tf32 on sm_80+, fp16 on sm_70/75).
// ---------------------------------------------------------------------------

namespace {

constexpr int KTILE  = 48;   // 45 real k values per channel, padded to 48
constexpr int MTILE  = 64;   // all output channels
constexpr int NTILE  = 64;   // pixels per block
constexpr int NTHREADS = 256;

template <typename T> struct ConvTraits;

template <> struct ConvTraits<float> {  // tf32 path
  static constexpr int KSTEP = 8;
  using a_frag = nvcuda::wmma::fragment<nvcuda::wmma::matrix_a, 16, 16, 8,
                                        nvcuda::wmma::precision::tf32,
                                        nvcuda::wmma::row_major>;
  using b_frag = nvcuda::wmma::fragment<nvcuda::wmma::matrix_b, 16, 16, 8,
                                        nvcuda::wmma::precision::tf32,
                                        nvcuda::wmma::row_major>;
  static __device__ __forceinline__ float cvt(float v) { return v; }
  static __device__ __forceinline__ void convert_a(a_frag &) {}
  static __device__ __forceinline__ void convert_b(b_frag &) {}
};

template <> struct ConvTraits<__half> {  // fp16 path
  static constexpr int KSTEP = 16;
  using a_frag = nvcuda::wmma::fragment<nvcuda::wmma::matrix_a, 16, 16, 16,
                                        __half, nvcuda::wmma::row_major>;
  using b_frag = nvcuda::wmma::fragment<nvcuda::wmma::matrix_b, 16, 16, 16,
                                        __half, nvcuda::wmma::row_major>;
  static __device__ __forceinline__ __half cvt(float v) { return __float2half_rn(v); }
  static __device__ __forceinline__ void convert_a(a_frag &) {}
  static __device__ __forceinline__ void convert_b(b_frag &) {}
};

template <typename T, bool TF32>
__global__ __launch_bounds__(NTHREADS)
void conv_wmma_kernel(const float* __restrict__ x,
                      const float* __restrict__ w,
                      const float* __restrict__ bias,
                      float* __restrict__ y,
                      int N, int C, int H, int W,
                      int OC, int OH, int OW, int KH, int KW) {
#if (__CUDA_ARCH__ >= 700) && (!TF32 || (__CUDA_ARCH__ >= 800))
  using Tr = ConvTraits<T>;
  constexpr int KSTEP = Tr::KSTEP;
  constexpr int NQ    = KTILE / KSTEP;

  // 24 KB of shared memory, reused for the output staging as well.
  __shared__ float smem_raw[6144];
  T* As = reinterpret_cast<T*>(smem_raw);                       // 64 x 48
  T* Bs = reinterpret_cast<T*>(smem_raw) + MTILE * KTILE;       // 48 x 64
  float* out_s = smem_raw;                                      // 64 x 64 (reuse)

  const int ow0   = blockIdx.x * NTILE;
  const int oh    = blockIdx.y;
  const int nb    = blockIdx.z;
  const int tid   = threadIdx.x;
  const int warp  = tid >> 5;
  const int mt    = warp & 3;          // 4 M-tiles of 16
  const int nt0   = (warp >> 2) * 2;   // each warp owns 2 N-tiles

  const int KHW = KH * KW;             // 45

  typename Tr::a_frag afr;
  typename Tr::b_frag bfr[2];
  nvcuda::wmma::fragment<nvcuda::wmma::accumulator, 16, 16, KSTEP, float> acc[2];
  nvcuda::wmma::fill_fragment(acc[0], 0.0f);
  nvcuda::wmma::fill_fragment(acc[1], 0.0f);

  const float* xb = x + (long)nb * C * H * W;

  for (int c = 0; c < C; ++c) {
    // ---- stage weights: As[oc][k] = w[oc][c][k], zero padded to 48 ----
    {
      const float* wc = w + (long)c * KHW;   // stride between oc rows is C*KHW
      #pragma unroll
      for (int it = 0; it < (MTILE * KTILE) / NTHREADS; ++it) {
        int idx = tid + it * NTHREADS;
        int m  = idx / KTILE;
        int k  = idx - m * KTILE;
        float v = 0.f;
        if (m < OC && k < KHW) v = wc[(long)m * (C * KHW) + k];
        As[idx] = Tr::cvt(v);
      }
    }
    // ---- stage input patch: Bs[k][n] = x[nb][c][oh+kh][ow0+n+kw] ----
    {
      #pragma unroll
      for (int it = 0; it < (KTILE * NTILE) / NTHREADS; ++it) {
        int idx = tid + it * NTHREADS;
        int k  = idx >> 6;
        int nn = idx & 63;
        float v = 0.f;
        if (k < KHW) {
          int kh = k / KW;
          int kw = k - kh * KW;
          int h  = oh + kh;
          int wc2 = ow0 + nn + kw;
          if (h < H && wc2 < W)
            v = xb[((long)c * H + h) * W + wc2];
        }
        Bs[idx] = Tr::cvt(v);
      }
    }
    __syncthreads();

    // ---- tensor-core accumulation over the 48 k values ----
    #pragma unroll
    for (int q = 0; q < NQ; ++q) {
      nvcuda::wmma::load_matrix_sync(afr, As + (mt * 16) * KTILE + q * KSTEP, KTILE);
      Tr::convert_a(afr);
      nvcuda::wmma::load_matrix_sync(bfr[0], Bs + (q * KSTEP) * NTILE + nt0 * 16, NTILE);
      nvcuda::wmma::load_matrix_sync(bfr[1], Bs + (q * KSTEP) * NTILE + (nt0 + 1) * 16, NTILE);
      Tr::convert_b(bfr[0]);
      Tr::convert_b(bfr[1]);
      nvcuda::wmma::mma_sync(acc[0], afr, bfr[0], acc[0]);
      nvcuda::wmma::mma_sync(acc[1], afr, bfr[1], acc[1]);
    }
    __syncthreads();
  }

  // ---- store accumulators to shared, then guarded copy to global ----
  nvcuda::wmma::store_matrix_sync(out_s + (mt * 16) * NTILE + nt0 * 16,
                                  acc[0], NTILE, nvcuda::wmma::mem_row_major);
  nvcuda::wmma::store_matrix_sync(out_s + (mt * 16) * NTILE + (nt0 + 1) * 16,
                                  acc[1], NTILE, nvcuda::wmma::mem_row_major);
  __syncthreads();

  #pragma unroll
  for (int it = 0; it < (MTILE * NTILE) / NTHREADS; ++it) {
    int idx = tid + it * NTHREADS;
    int m  = idx >> 6;
    int nn = idx & 63;
    int ow = ow0 + nn;
    if (m < OC && ow < OW) {
      float v = out_s[idx];
      if (bias != nullptr) v += __ldg(bias + m);
      y[(((long)nb * OC + m) * OH + oh) * OW + ow] = v;
    }
  }
#endif
}

}  // namespace

torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor weight,
                             c10::optional<torch::Tensor> bias_opt,
                             int64_t stride, int64_t padding,
                             int64_t dilation, int64_t groups) {
  const at::cuda::CUDAGuard device_guard(x.device());
  auto xc = x.contiguous();
  auto wc = weight.contiguous();
  torch::Tensor bc;
  const float* bias_ptr = nullptr;
  if (bias_opt.has_value() && bias_opt->defined()) {
    bc = bias_opt->contiguous();
    bias_ptr = bc.data_ptr<float>();
  }

  const int N  = xc.size(0);
  const int C  = xc.size(1);
  const int H  = xc.size(2);
  const int W  = xc.size(3);
  const int OC = wc.size(0);
  const int KH = wc.size(2);
  const int KW = wc.size(3);
  const int OH = (H + 2 * (int)padding - (int)dilation * (KH - 1) - 1) / (int)stride + 1;
  const int OW = (W + 2 * (int)padding - (int)dilation * (KW - 1) - 1) / (int)stride + 1;

  auto y = at::empty({N, OC, OH, OW}, xc.options());

  dim3 block(NTHREADS);
  dim3 grid((OW + NTILE - 1) / NTILE, OH, N);
  auto stream = at::cuda::getCurrentCUDAStream();

  int major = at::cuda::getCurrentDeviceProperties()->major;
  if (major >= 8) {
    conv_wmma_kernel<float, true><<<grid, block, 0, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
        y.data_ptr<float>(), N, C, H, W, OC, OH, OW, KH, KW);
  } else {
    conv_wmma_kernel<__half, false><<<grid, block, 0, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
        y.data_ptr<float>(), N, C, H, W, OC, OH, OW, KH, KW);
  }
  return y;
}
"""

conv_cpp_source = (
    "torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor weight, "
    "c10::optional<torch::Tensor> bias_opt, int64_t stride, int64_t padding, "
    "int64_t dilation, int64_t groups);"
)

conv_ext = load_inline(
    name="conv2d_tensorcore_ext",
    cpp_sources=conv_cpp_source,
    cuda_sources=conv_cuda_source,
    functions=["conv2d_forward"],
    verbose=False,
    extra_cflags=["-O3", "-std=c++17"],
    extra_cuda_cflags=["-O3", "-std=c++17"],
)


class ModelNew(nn.Module):
    """
    2D convolution accelerated with a custom implicit-GEMM tensor-core CUDA
    kernel (WMMA tf32 on sm_80+, WMMA fp16 on sm_70/75). Falls back to the
    standard PyTorch/cuDNN convolution for any configuration outside the
    fast path (or on unsupported hardware), so the module is fully general.
    """

    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple,
                 stride: int = 1, padding: int = 0, dilation: int = 1,
                 groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size,
                                stride=stride, padding=padding,
                                dilation=dilation, groups=groups, bias=bias)
        self._ext = conv_ext

    def _fast_path_ok(self, x: torch.Tensor) -> bool:
        cv = self.conv2d
        if not (x.is_cuda and cv.weight.is_cuda):
            return False
        if x.dtype != torch.float32 or cv.weight.dtype != torch.float32:
            return False
        if cv.groups != 1:
            return False
        if cv.stride != (1, 1) or cv.padding != (0, 0) or cv.dilation != (1, 1):
            return False
        kh, kw = cv.kernel_size
        if not (cv.in_channels == 32 and cv.out_channels == 64 and kh == 5 and kw == 9):
            return False
        try:
            if torch.cuda.get_device_capability(x.device)[0] < 7:
                return False
        except Exception:
            return False
        return True

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        cv = self.conv2d
        if self._fast_path_ok(x):
            return self._ext.conv2d_forward(
                x, cv.weight, cv.bias,
                cv.stride[0], cv.padding[0], cv.dilation[0], cv.groups)
        return cv(x)
```