```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

cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

#define TILE_M 128
#define TILE_N 64
#define TILE_K 32
#define NTHREADS 256
#define K_CHUNKS 8
#define LDA (TILE_M + 8)
#define LDB (TILE_N + 8)

__device__ __forceinline__ uint32_t f2tf32(float x) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    uint32_t u;
    asm volatile("cvt.rna.tf32.f32 %0, %1;" : "=r"(u) : "f"(x));
    return u;
#else
    return __float_as_uint(x);
#endif
}

__device__ __forceinline__ void mma16808_tf32(float* d, const uint32_t* a, const uint32_t* b) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    asm volatile(
        "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
        : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]));
#endif
}

// Rearrange ConvTranspose2d weight (CI,CO,4,4) into per-phase GEMM-friendly
// layout Wp[phase][k=(ci,a,b)][co] with kh=2a+1-po, kw=2b+1-pq.
__global__ void rearrange_weight_kernel(const float* __restrict__ w,
                                        float* __restrict__ wp) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= 4 * 256 * 64) return;
    int co = idx & 63;
    int k  = (idx >> 6) & 255;
    int p  = idx >> 14;
    int po = p >> 1, pq = p & 1;
    int ci = k >> 2;
    int a  = (k >> 1) & 1;
    int b  = k & 1;
    int kh = 2 * a + 1 - po;
    int kw = 2 * b + 1 - pq;
    wp[idx] = w[(long)ci * (64 * 16) + co * 16 + kh * 4 + kw];
}

// Implicit-GEMM transposed convolution (k=4, s=2, p=1, op=1) using TF32
// tensor cores (mma.sync PTX), with fused (+conv_bias, -bias, tanh) epilogue.
__global__ void deconv_gemm_kernel(const float* __restrict__ x,
                                   const float* __restrict__ wp,
                                   const float* __restrict__ conv_bias,
                                   const float* __restrict__ sub_bias,
                                   float* __restrict__ y,
                                   int N, int H, int W, int Ho, int Wo,
                                   int has_cb) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    const int phase = blockIdx.z;
    const int po = phase >> 1, pq = phase & 1;
    const int Hp = (Ho - po + 1) >> 1;
    const int Wp2 = (Wo - pq + 1) >> 1;
    const long M = (long)N * Hp * Wp2;
    const long row0 = (long)blockIdx.x * TILE_M;
    if (row0 >= M) return;

    extern __shared__ uint32_t smem[];
    uint32_t* As = smem;                              // [2][TILE_K][LDA]
    uint32_t* Bs = As + 2 * TILE_K * LDA;             // [2][TILE_K][LDB]
    int* soff  = (int*)(Bs + 2 * TILE_K * LDB);       // [TILE_M*4]
    long long* sbase = (long long*)(soff + TILE_M * 4); // [TILE_M]
    long long* sorow = sbase + TILE_M;                // [TILE_M]

    const int tid = threadIdx.x;
    const long xsc = (long)H * W;
    const long xsn = 64 * xsc;
    const long ytc = (long)Ho * Wo;
    const long ytn = 64 * ytc;

    // Per-row geometry precompute
    if (tid < TILE_M) {
        long row = row0 + tid;
        if (row < M) {
            long hw = (long)Hp * Wp2;
            long n = row / hw;
            long rem = row - n * hw;
            int oh2 = (int)(rem / Wp2);
            int ow2 = (int)(rem - (long)oh2 * Wp2);
            sbase[tid] = n * xsn;
            #pragma unroll
            for (int ab = 0; ab < 4; ++ab) {
                int a = ab >> 1, b = ab & 1;
                int ih = oh2 + po - a;
                int iw = ow2 + pq - b;
                bool v = (ih >= 0 && ih < H && iw >= 0 && iw < W);
                soff[tid * 4 + ab] = v ? (ih * W + iw) : -1;
            }
            int oh = 2 * oh2 + po;
            int ow = 2 * ow2 + pq;
            sorow[tid] = n * ytn + (long)oh * Wo + ow;
        } else {
            sbase[tid] = 0;
            #pragma unroll
            for (int ab = 0; ab < 4; ++ab) soff[tid * 4 + ab] = -1;
            sorow[tid] = -1;
        }
    }
    __syncthreads();

    const int wid  = tid >> 5;
    const int lane = tid & 31;
    const int wm = (wid & 3) * 32;
    const int wn = (wid >> 2) * 32;
    const int gid = lane >> 2;   // 0..7
    const int tig = lane & 3;    // 0..3

    float acc[2][4][4];
    #pragma unroll
    for (int i = 0; i < 2; ++i)
        #pragma unroll
        for (int j = 0; j < 4; ++j)
            #pragma unroll
            for (int r = 0; r < 4; ++r)
                acc[i][j][r] = 0.0f;

    const float* wphase = wp + phase * (256 * 64);

    // ---- prologue: load chunk 0 into buffer 0 ----
    {
        uint32_t* Asb = As;
        #pragma unroll
        for (int it = 0; it < (TILE_M * TILE_K) / NTHREADS; ++it) {
            int idx = tid + it * NTHREADS;
            int m = idx & (TILE_M - 1);
            int kk = idx >> 7;
            int ci = kk >> 2;
            int ab = kk & 3;
            int off = soff[m * 4 + ab];
            float v = 0.f;
            if (off >= 0) v = __ldg(x + sbase[m] + (long)ci * xsc + off);
            Asb[kk * LDA + m] = f2tf32(v);
        }
        uint32_t* Bsb = Bs;
        #pragma unroll
        for (int it = 0; it < (TILE_K * TILE_N) / NTHREADS; ++it) {
            int idx = tid + it * NTHREADS;
            int kk = idx >> 6;
            int n  = idx & 63;
            Bsb[kk * LDB + n] = f2tf32(__ldg(wphase + idx));
        }
    }
    __syncthreads();

    // ---- main K loop with double buffering ----
    #pragma unroll
    for (int kc = 0; kc < K_CHUNKS; ++kc) {
        int buf = kc & 1;
        if (kc + 1 < K_CHUNKS) {
            int nb = buf ^ 1;
            uint32_t* Asb = As + nb * (TILE_K * LDA);
            #pragma unroll
            for (int it = 0; it < (TILE_M * TILE_K) / NTHREADS; ++it) {
                int idx = tid + it * NTHREADS;
                int m = idx & (TILE_M - 1);
                int kk = idx >> 7;
                int ci = kk >> 2;
                int ab = kk & 3;
                int off = soff[m * 4 + ab];
                float v = 0.f;
                if (off >= 0) v = __ldg(x + sbase[m] + (long)ci * xsc + off);
                Asb[kk * LDA + m] = f2tf32(v);
            }
            uint32_t* Bsb = Bs + nb * (TILE_K * LDB);
            const float* wpp = wphase + (kc + 1) * (TILE_K * TILE_N);
            #pragma unroll
            for (int it = 0; it < (TILE_K * TILE_N) / NTHREADS; ++it) {
                int idx = tid + it * NTHREADS;
                int kk = idx >> 6;
                int n  = idx & 63;
                Bsb[kk * LDB + n] = f2tf32(__ldg(wpp + idx));
            }
        }
        {
            const uint32_t* Asb = As + buf * (TILE_K * LDA);
            const uint32_t* Bsb = Bs + buf * (TILE_K * LDB);
            #pragma unroll
            for (int ks = 0; ks < TILE_K / 8; ++ks) {
                const int k8 = ks * 8;
                uint32_t af[2][4];
                uint32_t bf[4][2];
                #pragma unroll
                for (int i = 0; i < 2; ++i) {
                    const int mbase = wm + 16 * i + gid;
                    af[i][0] = Asb[(k8 + tig)     * LDA + mbase];
                    af[i][1] = Asb[(k8 + tig)     * LDA + mbase + 8];
                    af[i][2] = Asb[(k8 + tig + 4) * LDA + mbase];
                    af[i][3] = Asb[(k8 + tig + 4) * LDA + mbase + 8];
                }
                #pragma unroll
                for (int j = 0; j < 4; ++j) {
                    const int nbase = wn + 8 * j + gid;
                    bf[j][0] = Bsb[(k8 + tig)     * LDB + nbase];
                    bf[j][1] = Bsb[(k8 + tig + 4) * LDB + nbase];
                }
                #pragma unroll
                for (int i = 0; i < 2; ++i)
                    #pragma unroll
                    for (int j = 0; j < 4; ++j)
                        mma16808_tf32(acc[i][j], af[i], bf[j]);
            }
        }
        __syncthreads();
    }

    // ---- epilogue: fused +conv_bias, -bias, tanh, direct global writes ----
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int h = 0; h < 2; ++h) {
            int row = wm + 16 * i + 8 * h + gid;
            long long orow = sorow[row];
            if (orow < 0) continue;
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                int co = wn + 8 * j + 2 * tig;
                float v0 = acc[i][j][2 * h];
                float v1 = acc[i][j][2 * h + 1];
                if (has_cb) {
                    v0 += __ldg(conv_bias + co);
                    v1 += __ldg(conv_bias + co + 1);
                }
                v0 -= __ldg(sub_bias + co);
                v1 -= __ldg(sub_bias + co + 1);
                y[orow + (long long)co * ytc] = tanhf(v0);
                y[orow + (long long)(co + 1) * ytc] = tanhf(v1);
            }
        }
    }
#endif
}

// Fused elementwise fallback: y = tanh(x + conv_bias - sub_bias) per channel
__global__ void fused_bias_tanh_kernel(const float* __restrict__ x,
                                       const float* __restrict__ conv_bias,
                                       const float* __restrict__ sub_bias,
                                       float* __restrict__ y,
                                       long HW, long total, int CO, int has_cb) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int co = (int)((idx / HW) % CO);
    float v = x[idx];
    if (has_cb) v += __ldg(conv_bias + co);
    v -= __ldg(sub_bias + co);
    y[idx] = tanhf(v);
}

static bool g_smem_configured = false;

torch::Tensor deconv_forward(torch::Tensor x, torch::Tensor weight,
                             c10::optional<torch::Tensor> conv_bias,
                             torch::Tensor sub_bias,
                             int64_t stride, int64_t padding,
                             int64_t output_padding) {
    TORCH_CHECK(x.is_cuda() && x.dtype() == torch::kFloat32, "x must be CUDA fp32");
    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = sub_bias.contiguous();

    const int N = xc.size(0);
    const int CI = xc.size(1);
    const int H = xc.size(2);
    const int W = xc.size(3);
    const int CO = wc.size(1);
    const int KH = wc.size(2);
    const int KW = wc.size(3);

    const long Ho = (long)(H - 1) * stride - 2 * padding + KH + output_padding;
    const long Wo = (long)(W - 1) * stride - 2 * padding + KW + output_padding;

    auto stream = at::cuda::getCurrentCUDAStream();

    torch::Tensor cb;
    const float* cb_ptr = nullptr;
    int has_cb = 0;
    if (conv_bias.has_value() && conv_bias->defined()) {
        cb = conv_bias->contiguous();
        cb_ptr = cb.data_ptr<float>();
        has_cb = 1;
    }

    // check fast-path feasibility
    int dev = 0; cudaGetDevice(&dev);
    static int arch_major = -1;
    if (arch_major < 0) {
        cudaDeviceGetAttribute(&arch_major, cudaDevAttrComputeCapabilityMajor, dev);
    }
    bool fast = (arch_major == 8) && CI == 64 && CO == 64 && KH == 4 && KW == 4 &&
                stride == 2 && padding == 1 && output_padding == 1 &&
                Ho <= INT_MAX && Wo <= INT_MAX;

    auto y = torch::empty({N, CO, Ho, Wo}, xc.options());

    if (!fast) {
        auto tmp = at::conv_transpose2d(xc, wc, conv_bias.has_value() ? cb : torch::Tensor(),
                                        {stride, stride}, {padding, padding},
                                        {output_padding, output_padding}, 1);
        long total = tmp.numel();
        long HW = Ho * Wo;
        int threads = 256;
        long blocks = (total + threads - 1) / threads;
        fused_bias_tanh_kernel<<<(unsigned)blocks, threads, 0, stream>>>(
            tmp.data_ptr<float>(), cb_ptr, bc.data_ptr<float>(),
            y.data_ptr<float>(), HW, total, CO, has_cb);
        return y;
    }

    // rearranged weights (4 phases x 256 x 64)
    auto wp = torch::empty({4 * 256 * 64}, xc.options());
    rearrange_weight_kernel<<<256, 256, 0, stream>>>(wc.data_ptr<float>(), wp.data_ptr<float>());

    // tiles per phase (max over phases)
    long maxM = 0;
    for (int p = 0; p < 4; ++p) {
        int po = p >> 1, pq = p & 1;
        long Hp = (Ho - po + 1) >> 1;
        long Wp2 = (Wo - pq + 1) >> 1;
        long Mp = (long)N * Hp * Wp2;
        if (Mp > maxM) maxM = Mp;
    }
    long tiles = (maxM + TILE_M - 1) / TILE_M;

    size_t smem = (size_t)(2 * TILE_K * LDA + 2 * TILE_K * LDB) * sizeof(uint32_t)
                + (size_t)(TILE_M * 4) * sizeof(int)
                + (size_t)(TILE_M * 2) * sizeof(long long);
    if (!g_smem_configured) {
        cudaFuncSetAttribute(deconv_gemm_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
        g_smem_configured = true;
    }
    dim3 grid((unsigned)tiles, 1, 4);
    deconv_gemm_kernel<<<grid, NTHREADS, smem, stream>>>(
        xc.data_ptr<float>(), wp.data_ptr<float>(), cb_ptr, bc.data_ptr<float>(),
        y.data_ptr<float>(), N, H, W, (int)Ho, (int)Wo, has_cb);
    return y;
}
"""

cpp_source = """
torch::Tensor deconv_forward(torch::Tensor x, torch::Tensor weight,
                             c10::optional<torch::Tensor> conv_bias,
                             torch::Tensor sub_bias,
                             int64_t stride, int64_t padding,
                             int64_t output_padding);
"""

deconv_mod = load_inline(
    name="deconv_fused_mma",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["deconv_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model: transposed convolution via a TF32 tensor-core implicit-GEMM
    CUDA kernel (mma.sync.m16n8k8) with fused (+conv_bias, -bias, tanh) epilogue.
    """
    def __init__(self, in_channels, out_channels, kernel_size, bias_shape, stride=2, padding=1, output_padding=1):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size,
                                                 stride=stride, padding=padding,
                                                 output_padding=output_padding)
        self.bias = nn.Parameter(torch.randn(bias_shape))
        self.stride = stride
        self.padding = padding
        self.output_padding = output_padding

    def forward(self, x):
        return deconv_mod.deconv_forward(
            x,
            self.conv_transpose.weight,
            self.conv_transpose.bias,
            self.bias.view(-1),
            self.stride,
            self.padding,
            self.output_padding,
        )
```