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

// ---------------------------------------------------------------------------
// Fused ConvTranspose2d(k=4,s=2,p=1,op=1) + bias-subtract + tanh.
//
// Parity-phase decomposition: for phase (po,pq) = (oh&1, ow&1) the op is an
// exact GEMM over output pixels M=(n,oh2,ow2), K=(ci,a,b), N=co:
//   out = sum_{ci,a,b} x[n,ci,oh2+po-a,ow2+pq-b] * w[ci,co,2a+1-po,2b+1-pq]
// with oh = 2*oh2+po, ow = 2*ow2+pq.  Works for arbitrary (even/odd) input
// and output spatial sizes; out-of-range input taps are zero-filled.
//
// TF32 tensor-core (mma.sync m16n8k8) implicit GEMM, 3-stage cp.async
// pipeline, fused (+conv_bias -bias, tanh) epilogue.
// gridDim.x = phase (4): the 4 phase-CTAs of a tile are scheduled together,
// so their interleaved output writes merge in L2 and input reads are shared.
// ---------------------------------------------------------------------------

#define TILE_M 128
#define TILE_N 64
#define TILE_K 32
#define KTILES 8
#define STAGES 3
#define NTHREADS 256
#define LDA 136
#define LDB 72
#define LDC 136

#define PIPE_SMEM ((STAGES * (TILE_K * LDA + TILE_K * LDB)) * 4)
#define CS_SMEM (TILE_N * LDC * 4)
#define DYN_SMEM (PIPE_SMEM > CS_SMEM ? PIPE_SMEM : CS_SMEM)

__device__ __forceinline__ unsigned smem_u32(const void* p) {
    return (unsigned)__cvta_generic_to_shared(p);
}

__device__ __forceinline__ void cp_async4(void* dst, const void* src, bool p) {
    int sz = p ? 4 : 0;
    asm volatile("cp.async.ca.shared.global [%0], [%1], 4, %2;\n" ::
                 "r"(smem_u32(dst)), "l"(src), "r"(sz));
}

__device__ __forceinline__ void cp_async16(void* dst, const void* src) {
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::
                 "r"(smem_u32(dst)), "l"(src));
}

__device__ __forceinline__ void cp_commit() {
    asm volatile("cp.async.commit_group;\n" ::);
}

template <int N>
__device__ __forceinline__ void cp_wait() {
    asm volatile("cp.async.wait_group %0;\n" :: "n"(N));
}

__device__ __forceinline__ float to_tf32(float x) {
    unsigned u = __float_as_uint(x);
    u = (u + 0x1000u) & ~0x1FFFu;
    return __uint_as_float(u);
}

__device__ __forceinline__ void mma_tf32(float* c, const float* a, const float* b) {
    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};\n"
        : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
        : "r"(__float_as_uint(a[0])), "r"(__float_as_uint(a[1])),
          "r"(__float_as_uint(a[2])), "r"(__float_as_uint(a[3])),
          "r"(__float_as_uint(b[0])), "r"(__float_as_uint(b[1])));
}

// wp[phase][k=(ci,a,b)][co] (tf32-rounded); comb[co] = conv_bias - sub_bias
__global__ void rearrange_weight_kernel(const float* __restrict__ w,
                                        const float* __restrict__ conv_bias,
                                        const float* __restrict__ sub_bias,
                                        float* __restrict__ wp,
                                        float* __restrict__ comb,
                                        int has_cb) {
    if (blockIdx.x == 256) {
        int co = threadIdx.x;
        if (co < 64) {
            float cb = has_cb ? __ldg(conv_bias + co) : 0.0f;
            comb[co] = cb - __ldg(sub_bias + co);
        }
        return;
    }
    int idx = blockIdx.x * blockDim.x + threadIdx.x;   // 4*256*64 elements
    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] = to_tf32(w[(long)ci * 1024 + co * 16 + kh * 4 + kw]);
}

__global__ void __launch_bounds__(NTHREADS) deconv_fused_tc_kernel(
        const float* __restrict__ x,
        const float* __restrict__ wp,
        const float* __restrict__ comb,
        float* __restrict__ y,
        int N, int H, int W, int Ho, int Wo) {

    extern __shared__ char dynsmem[];
    float* As = reinterpret_cast<float*>(dynsmem);                       // [STAGES][TILE_K][LDA]
    float* Bs = reinterpret_cast<float*>(dynsmem) + STAGES * TILE_K * LDA; // [STAGES][TILE_K][LDB]
    float* Cs = reinterpret_cast<float*>(dynsmem);                       // reused at epilogue

    __shared__ int soff[TILE_M][4];
    __shared__ long long sbase[TILE_M];
    __shared__ long long sorow[TILE_M];
    __shared__ float scomb[TILE_N];

    const int phase = blockIdx.x;
    const int po = phase >> 1, pq = phase & 1;
    const int Hp = (Ho - po + 1) >> 1;    // ceil((Ho-po)/2), valid for odd Ho
    const int Wp2 = (Wo - pq + 1) >> 1;
    const long M = (long)N * Hp * Wp2;
    const long row0 = (long)blockIdx.y * TILE_M;
    if (row0 >= M) return;   // whole-CTA uniform exit

    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;

    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] = (long long)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][ab] = v ? (ih * W + iw) : -1;
            }
            int oh = 2 * oh2 + po;
            int ow = 2 * ow2 + pq;
            sorow[tid] = (long long)n * ytn + (long)oh * Wo + ow;
        } else {
            sbase[tid] = 0;
            #pragma unroll
            for (int ab = 0; ab < 4; ++ab) soff[tid][ab] = -1;
            sorow[tid] = -1;
        }
    }
    if (tid < TILE_N) scomb[tid] = __ldg(comb + tid);
    __syncthreads();

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

    auto stage = [&](int ks) {
        float* Asb = As + (ks % STAGES) * (TILE_K * LDA);
        float* Bsb = Bs + (ks % STAGES) * (TILE_K * LDB);
        int kbase = ks * TILE_K;
        #pragma unroll
        for (int it = 0; it < (TILE_K * TILE_M) / NTHREADS; ++it) {  // 16
            int idx = tid + it * NTHREADS;
            int m  = idx & (TILE_M - 1);
            int kk = idx >> 7;
            int kg = kbase + kk;
            int ci = kg >> 2;
            int ab = kg & 3;
            int off = soff[m][ab];
            const float* src = x + (sbase[m] + (long)ci * xsc + (off >= 0 ? off : 0));
            cp_async4(&Asb[kk * LDA + m], src, off >= 0);
        }
        #pragma unroll
        for (int it = 0; it < (TILE_K * TILE_N) / (NTHREADS * 4); ++it) {  // 4
            int idx = tid + it * NTHREADS;   // 0..1023
            int co4 = (idx & 15) * 4;
            int kk  = idx >> 4;
            cp_async16(&Bsb[kk * LDB + co4], wphase + (long)(kbase + kk) * 64 + co4);
        }
        cp_commit();
    };

    #pragma unroll
    for (int s = 0; s < STAGES - 1; ++s) stage(s);

    const int warp = tid >> 5;
    const int lane = tid & 31;
    const int m_off = (warp & 3) * 32;
    const int n_off = (warp >> 2) * 32;
    const int gid = lane >> 2;
    const int tig = lane & 3;

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

    cp_wait<STAGES - 2>();
    __syncthreads();

    #pragma unroll
    for (int ks = 0; ks < KTILES; ++ks) {
        if (ks + STAGES - 1 < KTILES) stage(ks + STAGES - 1);

        const float* Asb = As + (ks % STAGES) * (TILE_K * LDA);
        const float* Bsb = Bs + (ks % STAGES) * (TILE_K * LDB);

        #pragma unroll
        for (int k8 = 0; k8 < TILE_K / 8; ++k8) {
            int kb = k8 * 8;
            float af[2][4], bf[4][2];
            #pragma unroll
            for (int e = 0; e < 4; ++e) {
                int kk = kb + tig + 4 * (e >> 1);
                int mm = gid + 8 * (e & 1);
                af[0][e] = to_tf32(Asb[kk * LDA + m_off + mm]);
                af[1][e] = to_tf32(Asb[kk * LDA + m_off + 16 + mm]);
            }
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                #pragma unroll
                for (int e = 0; e < 2; ++e) {
                    bf[j][e] = Bsb[(kb + tig + 4 * e) * LDB + n_off + j * 8 + gid];
                }
            }
            #pragma unroll
            for (int i = 0; i < 2; ++i)
                #pragma unroll
                for (int j = 0; j < 4; ++j)
                    mma_tf32(acc[i * 4 + j], af[i], bf[j]);
        }

        cp_wait<STAGES - 2>();
        __syncthreads();
    }

    // epilogue: store accumulators to smem as [co][m]
    __syncthreads();
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            float* a8 = acc[i * 4 + j];
            int m0 = m_off + i * 16 + gid;
            int c0 = n_off + j * 8 + 2 * tig;
            Cs[c0 * LDC + m0]         = a8[0];
            Cs[(c0 + 1) * LDC + m0]   = a8[1];
            Cs[c0 * LDC + m0 + 8]     = a8[2];
            Cs[(c0 + 1) * LDC + m0 + 8] = a8[3];
        }
    }
    __syncthreads();

    // fused (+conv_bias -bias) + tanh, write out (m fastest for coalescing;
    // phase CTAs are adjacent in gridDim.x so lines merge in L2)
    const int me = tid & (TILE_M - 1);
    const int cg = tid >> 7;   // 0 or 1
    const long long orow = sorow[me];
    if (orow >= 0) {
        #pragma unroll
        for (int it = 0; it < (TILE_M * TILE_N) / NTHREADS; ++it) {  // 32
            int co = 2 * it + cg;
            float v = Cs[co * LDC + me] + scomb[co];
            y[orow + (long long)co * ytc] = tanhf(v);
        }
    }
}

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

static bool g_smem_set = 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;
    }

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

    // The tensor-core kernel handles arbitrary (including odd) spatial sizes:
    // per-phase sizes use ceil division and taps are bounds-checked.
    bool fast = (CI == 64) && (CO == 64) && (KH == 4) && (KW == 4) &&
                (stride == 2) && (padding == 1) && (output_padding == 1) &&
                (Ho <= INT_MAX) && (Wo <= INT_MAX) && (H > 0) && (W > 0);

    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>(), bc.data_ptr<float>(), y.data_ptr<float>(),
            HW, total, CO);
        return y;
    }

    if (!g_smem_set) {
        cudaFuncSetAttribute(deconv_fused_tc_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, DYN_SMEM);
        g_smem_set = true;
    }

    auto wp = torch::empty({4 * 256 * 64}, xc.options());
    auto comb = torch::empty({64}, xc.options());
    rearrange_weight_kernel<<<257, 256, 0, stream>>>(
        wc.data_ptr<float>(), cb_ptr, bc.data_ptr<float>(),
        wp.data_ptr<float>(), comb.data_ptr<float>(), has_cb);

    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;

    dim3 grid(4, (unsigned)tiles, 1);
    deconv_fused_tc_kernel<<<grid, NTHREADS, DYN_SMEM, stream>>>(
        xc.data_ptr<float>(), wp.data_ptr<float>(), comb.data_ptr<float>(),
        y.data_ptr<float>(), N, H, W, (int)Ho, (int)Wo);
    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_tf32_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 fused TF32 tensor-core
    implicit-GEMM CUDA kernel (mma.sync m16n8k8, 3-stage cp.async pipeline)
    with (+conv_bias, -bias, tanh) fused in the epilogue.  Parity-phase
    decomposition turns the k=4/s=2 deconv into exact GEMMs; the 4 phase
    CTAs of a tile are adjacent in the launch grid so interleaved writes
    merge in L2 and input reads are shared.  Handles arbitrary (even/odd)
    spatial sizes; falls back to ATen conv_transpose2d + a fused
    elementwise kernel for other configurations.
    """
    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,
        )
```