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

# ---------------------------------------------------------------------------
# Custom fused implementation of:
#   ConvTranspose2d(stride=2, padding=1, kernel=4, output_padding=1)
#   + per-channel bias subtraction
#   + tanh activation
#
# Strategy:
#  * A strided transposed convolution with k=4, s=2, p=1 is decomposed into
#    4 "phases" (row parity p, col parity q of the output pixel).  Each phase
#    is a dense 2x2 strided gather from the input -> an implicit GEMM with
#    M = batch*out_pixels, K = in_channels*4, N = out_channels.
#  * One kernel block processes a tile of 64 consecutive "super-pixels" of a
#    fixed row-parity p and computes BOTH column parities q (so the final
#    stores are fully contiguous) for all out_channels.
#  * Register blocking: 256 threads/block, each thread holds 4(m) x 4(oc) x
#    2(q) = 32 float accumulators.  Double-buffered shared-memory pipeline.
#  * Bias add (conv bias), bias subtraction (the model's bias parameter) and
#    tanh are fused into the epilogue -> no extra global memory passes.
# ---------------------------------------------------------------------------

cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

#define BM 64          // super-pixels per block (flattened m dimension)
#define BN 64          // out channels per block (all of them)
#define BK 16          // k-chunk
#define NTHREADS 256

// ---------------------------------------------------------------------------
// Weight reorganisation kernel.
// Wbuf[p][k'][qq][oc] = weight[oc][ic][3-p-2*ta][3-qq-2*tb]
//   with k' = ic*4 + ta*2 + tb
// ---------------------------------------------------------------------------
__global__ void build_wbuf_kernel(const float* __restrict__ w,
                                  float* __restrict__ wbuf,
                                  int IC, int OC) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = 2 * (IC * 4) * 2 * OC;
    if (idx >= total) return;

    int oc  = idx % OC;
    int qq  = (idx / OC) % 2;
    int kp  = (idx / (OC * 2)) % (IC * 4);
    int p   = idx / (OC * 2 * IC * 4);

    int ic = kp >> 2;
    int ta = (kp >> 1) & 1;
    int tb = kp & 1;

    int kh = 3 - p - 2 * ta;
    int kw = 3 - qq - 2 * tb;

    wbuf[((p * (IC * 4) + kp) * 2 + qq) * OC + oc] =
        w[((oc * IC + ic) * 4 + kh) * 4 + kw];
}

// ---------------------------------------------------------------------------
// Fused conv-transpose + bias + tanh kernel.
//
// Output pixel (h = 2r + p, w = 2c + q):
//   out = sum_{ic,ta,tb} x[n][ic][r+p-1+ta][c+q-1+tb]
//                        * w[oc][ic][3-p-2ta][3-q-2tb]
// Grid: (ceil(Mp/BM), OC/BN, n*2+p), block = 256 threads.
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(NTHREADS)
fused_tconv_bias_tanh_kernel(const float* __restrict__ x,
                             const float* __restrict__ wbuf,
                             const float* __restrict__ cbias,
                             const float* __restrict__ subbias,
                             float* __restrict__ out,
                             int batch, int IC, int H, int W,
                             int OC, int OH, int OW,
                             int Rp, int Cp, long Mp) {
    const int p  = blockIdx.z & 1;
    const int n  = blockIdx.z >> 1;
    const int oc0 = blockIdx.y * BN;
    const long m0 = (long)blockIdx.x * BM;

    __shared__ float As[2][2][BK][BM];   // [buf][q][k'][m]
    __shared__ float Bs[2][2][BK][BN];   // [buf][q][k'][oc]

    const int tid = threadIdx.x;
    const int tm  = tid % 16;            // m dimension (4 m's per thread)
    const int tn  = tid / 16;            // oc dimension (4 oc's per thread)

    // ---- precompute this thread's A-load geometry ------------------------
    const int mi_ld   = tid % 64;              // fixed for all 8 loads
    const long m_ld   = m0 + mi_ld;
    const bool m_ok   = (m_ld < Mp);
    const long rc     = m_ok ? (m_ld / Cp) : 0;
    const long cc     = m_ok ? (m_ld - rc * Cp) : 0;
    const int r_ld    = (int)rc;
    const int c_ld    = (int)cc;
    const int rowbase = r_ld + p - 1;          // + ta
    const int colbase = c_ld - 1;              // + tb + q
    const long xstride_c = (long)W;
    const long xbase  = ((long)n * IC) * H * W;

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

    const float* wbufp = wbuf + (long)p * (IC * 4) * 2 * OC;

    // -------------------- shared-tile load helpers ------------------------
    // load chunk kc (k' in [kc*BK, kc*BK+BK)) into buffer `buf`
    auto load_tile = [&](int kc, int buf) {
        float av[8], bv[8];
        #pragma unroll
        for (int pass = 0; pass < 8; ++pass) {
            int e   = tid + pass * NTHREADS;   // 0..2047
            int arr = e >> 10;                 // q
            int rem = e & 1023;
            int kp  = rem >> 6;                // 0..15
            // A gather
            {
                int kgl = kc * BK + kp;        // global k'
                int ic  = kgl >> 2;
                int ta  = (kgl >> 1) & 1;
                int tb  = kgl & 1;
                int row = rowbase + ta;
                int col = colbase + tb + arr;
                float v = 0.f;
                if (m_ok && row >= 0 && row < H && col >= 0 && col < W) {
                    v = x[xbase + (long)ic * H * W + (long)row * xstride_c + col];
                }
                av[pass] = v;
            }
            // B gather
            {
                int ocol = rem & 63;           // oc within tile
                int kgl  = kc * BK + kp;
                bv[pass] = wbufp[((long)kgl * 2 + arr) * OC + oc0 + ocol];
            }
        }
        #pragma unroll
        for (int pass = 0; pass < 8; ++pass) {
            int e   = tid + pass * NTHREADS;
            int arr = e >> 10;
            int rem = e & 1023;
            int kp  = rem >> 6;
            int ocol = rem & 63;
            As[buf][arr][kp][mi_ld] = av[pass];
            Bs[buf][arr][kp][ocol]  = bv[pass];
        }
    };

    const int kchunks = (IC * 4) / BK;         // exact multiple

    load_tile(0, 0);
    __syncthreads();

    for (int kc = 0; kc < kchunks; ++kc) {
        int buf = kc & 1;
        if (kc + 1 < kchunks) {
            // prefetch next chunk into registers, then shared
            float av[8], bv[8];
            int kn = kc + 1;
            #pragma unroll
            for (int pass = 0; pass < 8; ++pass) {
                int e   = tid + pass * NTHREADS;
                int arr = e >> 10;
                int rem = e & 1023;
                int kp  = rem >> 6;
                {
                    int kgl = kn * BK + kp;
                    int ic  = kgl >> 2;
                    int ta  = (kgl >> 1) & 1;
                    int tb  = kgl & 1;
                    int row = rowbase + ta;
                    int col = colbase + tb + arr;
                    float v = 0.f;
                    if (m_ok && row >= 0 && row < H && col >= 0 && col < W) {
                        v = x[xbase + (long)ic * H * W + (long)row * xstride_c + col];
                    }
                    av[pass] = v;
                }
                {
                    int ocol = rem & 63;
                    int kgl  = kn * BK + kp;
                    bv[pass] = wbufp[((long)kgl * 2 + arr) * OC + oc0 + ocol];
                }
            }
            // compute on current buffer
            #pragma unroll
            for (int kk = 0; kk < BK; ++kk) {
                float4 a0 = *reinterpret_cast<const float4*>(&As[buf][0][kk][tm * 4]);
                float4 a1 = *reinterpret_cast<const float4*>(&As[buf][1][kk][tm * 4]);
                float4 b0 = *reinterpret_cast<const float4*>(&Bs[buf][0][kk][tn * 4]);
                float4 b1 = *reinterpret_cast<const float4*>(&Bs[buf][1][kk][tn * 4]);
                acc[0][0][0] = fmaf(a0.x, b0.x, acc[0][0][0]);
                acc[0][0][1] = fmaf(a0.x, b0.y, acc[0][0][1]);
                acc[0][0][2] = fmaf(a0.x, b0.z, acc[0][0][2]);
                acc[0][0][3] = fmaf(a0.x, b0.w, acc[0][0][3]);
                acc[0][1][0] = fmaf(a0.y, b0.x, acc[0][1][0]);
                acc[0][1][1] = fmaf(a0.y, b0.y, acc[0][1][1]);
                acc[0][1][2] = fmaf(a0.y, b0.z, acc[0][1][2]);
                acc[0][1][3] = fmaf(a0.y, b0.w, acc[0][1][3]);
                acc[0][2][0] = fmaf(a0.z, b0.x, acc[0][2][0]);
                acc[0][2][1] = fmaf(a0.z, b0.y, acc[0][2][1]);
                acc[0][2][2] = fmaf(a0.z, b0.z, acc[0][2][2]);
                acc[0][2][3] = fmaf(a0.z, b0.w, acc[0][2][3]);
                acc[0][3][0] = fmaf(a0.w, b0.x, acc[0][3][0]);
                acc[0][3][1] = fmaf(a0.w, b0.y, acc[0][3][1]);
                acc[0][3][2] = fmaf(a0.w, b0.z, acc[0][3][2]);
                acc[0][3][3] = fmaf(a0.w, b0.w, acc[0][3][3]);
                acc[1][0][0] = fmaf(a1.x, b1.x, acc[1][0][0]);
                acc[1][0][1] = fmaf(a1.x, b1.y, acc[1][0][1]);
                acc[1][0][2] = fmaf(a1.x, b1.z, acc[1][0][2]);
                acc[1][0][3] = fmaf(a1.x, b1.w, acc[1][0][3]);
                acc[1][1][0] = fmaf(a1.y, b1.x, acc[1][1][0]);
                acc[1][1][1] = fmaf(a1.y, b1.y, acc[1][1][1]);
                acc[1][1][2] = fmaf(a1.y, b1.z, acc[1][1][2]);
                acc[1][1][3] = fmaf(a1.y, b1.w, acc[1][1][3]);
                acc[1][2][0] = fmaf(a1.z, b1.x, acc[1][2][0]);
                acc[1][2][1] = fmaf(a1.z, b1.y, acc[1][2][1]);
                acc[1][2][2] = fmaf(a1.z, b1.z, acc[1][2][2]);
                acc[1][2][3] = fmaf(a1.z, b1.w, acc[1][2][3]);
                acc[1][3][0] = fmaf(a1.w, b1.x, acc[1][3][0]);
                acc[1][3][1] = fmaf(a1.w, b1.y, acc[1][3][1]);
                acc[1][3][2] = fmaf(a1.w, b1.z, acc[1][3][2]);
                acc[1][3][3] = fmaf(a1.w, b1.w, acc[1][3][3]);
            }
            __syncthreads();
            #pragma unroll
            for (int pass = 0; pass < 8; ++pass) {
                int e   = tid + pass * NTHREADS;
                int arr = e >> 10;
                int rem = e & 1023;
                int kp  = rem >> 6;
                int ocol = rem & 63;
                As[buf ^ 1][arr][kp][mi_ld] = av[pass];
                Bs[buf ^ 1][arr][kp][ocol]  = bv[pass];
            }
            __syncthreads();
        } else {
            // last chunk, no prefetch
            #pragma unroll
            for (int kk = 0; kk < BK; ++kk) {
                float4 a0 = *reinterpret_cast<const float4*>(&As[buf][0][kk][tm * 4]);
                float4 a1 = *reinterpret_cast<const float4*>(&As[buf][1][kk][tm * 4]);
                float4 b0 = *reinterpret_cast<const float4*>(&Bs[buf][0][kk][tn * 4]);
                float4 b1 = *reinterpret_cast<const float4*>(&Bs[buf][1][kk][tn * 4]);
                acc[0][0][0] = fmaf(a0.x, b0.x, acc[0][0][0]);
                acc[0][0][1] = fmaf(a0.x, b0.y, acc[0][0][1]);
                acc[0][0][2] = fmaf(a0.x, b0.z, acc[0][0][2]);
                acc[0][0][3] = fmaf(a0.x, b0.w, acc[0][0][3]);
                acc[0][1][0] = fmaf(a0.y, b0.x, acc[0][1][0]);
                acc[0][1][1] = fmaf(a0.y, b0.y, acc[0][1][1]);
                acc[0][1][2] = fmaf(a0.y, b0.z, acc[0][1][2]);
                acc[0][1][3] = fmaf(a0.y, b0.w, acc[0][1][3]);
                acc[0][2][0] = fmaf(a0.z, b0.x, acc[0][2][0]);
                acc[0][2][1] = fmaf(a0.z, b0.y, acc[0][2][1]);
                acc[0][2][2] = fmaf(a0.z, b0.z, acc[0][2][2]);
                acc[0][2][3] = fmaf(a0.z, b0.w, acc[0][2][3]);
                acc[0][3][0] = fmaf(a0.w, b0.x, acc[0][3][0]);
                acc[0][3][1] = fmaf(a0.w, b0.y, acc[0][3][1]);
                acc[0][3][2] = fmaf(a0.w, b0.z, acc[0][3][2]);
                acc[0][3][3] = fmaf(a0.w, b0.w, acc[0][3][3]);
                acc[1][0][0] = fmaf(a1.x, b1.x, acc[1][0][0]);
                acc[1][0][1] = fmaf(a1.x, b1.y, acc[1][0][1]);
                acc[1][0][2] = fmaf(a1.x, b1.z, acc[1][0][2]);
                acc[1][0][3] = fmaf(a1.x, b1.w, acc[1][0][3]);
                acc[1][1][0] = fmaf(a1.y, b1.x, acc[1][1][0]);
                acc[1][1][1] = fmaf(a1.y, b1.y, acc[1][1][1]);
                acc[1][1][2] = fmaf(a1.y, b1.z, acc[1][1][2]);
                acc[1][1][3] = fmaf(a1.y, b1.w, acc[1][1][3]);
                acc[1][2][0] = fmaf(a1.z, b1.x, acc[1][2][0]);
                acc[1][2][1] = fmaf(a1.z, b1.y, acc[1][2][1]);
                acc[1][2][2] = fmaf(a1.z, b1.z, acc[1][2][2]);
                acc[1][2][3] = fmaf(a1.z, b1.w, acc[1][2][3]);
                acc[1][3][0] = fmaf(a1.w, b1.x, acc[1][3][0]);
                acc[1][3][1] = fmaf(a1.w, b1.y, acc[1][3][1]);
                acc[1][3][2] = fmaf(a1.w, b1.z, acc[1][3][2]);
                acc[1][3][3] = fmaf(a1.w, b1.w, acc[1][3][3]);
            }
        }
    }

    // --------------------------- epilogue ---------------------------------
    const long ostride = (long)OH * OW;
    float bnet[4];
    #pragma unroll
    for (int j = 0; j < 4; ++j) {
        int oc = oc0 + tn * 4 + j;
        bnet[j] = cbias[oc] - subbias[oc];
    }

    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        long m = m0 + tm * 4 + i;
        if (m >= Mp) continue;
        long r = m / Cp;
        long c = m - r * Cp;
        long h = 2 * r + p;
        if (h >= OH) continue;
        long w0 = 2 * c;
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            int oc = oc0 + tn * 4 + j;
            float* optr = out + ((long)n * OC + oc) * ostride + h * OW + w0;
            float v0 = tanhf(acc[0][i][j] + bnet[j]);
            optr[0] = v0;
            if (w0 + 1 < OW) {
                float v1 = tanhf(acc[1][i][j] + bnet[j]);
                optr[1] = v1;
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Host wrapper
// ---------------------------------------------------------------------------
torch::Tensor fused_tconv_bias_tanh(torch::Tensor x,
                                    torch::Tensor weight,
                                    torch::Tensor cbias,
                                    torch::Tensor subbias,
                                    int64_t stride, int64_t padding,
                                    int64_t output_padding) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "inputs must be CUDA");
    TORCH_CHECK(x.dim() == 4, "x must be NCHW");
    TORCH_CHECK(weight.size(2) == 4 && weight.size(3) == 4,
                "kernel specialised for 4x4");
    TORCH_CHECK(stride == 2 && padding == 1 && output_padding == 1,
                "kernel specialised for s=2,p=1,op=1");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto cc = cbias.contiguous();
    auto sc = subbias.contiguous();

    const long batch = xc.size(0);
    const long IC    = xc.size(1);
    const long H     = xc.size(2);
    const long W     = xc.size(3);
    const long OC    = wc.size(0);
    const long OH    = (H - 1) * 2 - 2 * 1 + 4 + 1;   // s=2,p=1,k=4,op=1
    const long OW    = (W - 1) * 2 - 2 * 1 + 4 + 1;

    auto out = torch::empty({batch, OC, OH, OW}, xc.options());
    auto wbuf = torch::empty({2 * IC * 4 * 2 * OC}, xc.options());

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

    // weight reorganisation
    {
        int total = 2 * (int)(IC * 4) * 2 * (int)OC;
        int threads = 256;
        int blocks = (total + threads - 1) / threads;
        build_wbuf_kernel<<<blocks, threads, 0, stream>>>(
            wc.data_ptr<float>(), wbuf.data_ptr<float>(), (int)IC, (int)OC);
    }

    const long R[2] = {(OH + 1) / 2, OH / 2};
    const long C[2] = {(OW + 1) / 2, OW / 2};
    const long Mp[2] = {batch * R[0] * C[0], batch * R[1] * C[1]};
    const long maxM = Mp[0] > Mp[1] ? Mp[0] : Mp[1];

    dim3 grid((unsigned)((maxM + BM - 1) / BM),
              (unsigned)(OC / BN),
              (unsigned)(batch * 2));

    // The two phases have slightly different M; use per-phase Cp/Mp via p.
    // We pass both and select inside through blockIdx.z parity using Rp/Cp/Mp
    // of the *maximum*; to keep it simple we relaunch per-phase-correct values
    // by choosing Rp/Cp/Mp of phase 0 and guarding, but exact per-phase dims
    // avoid wasted blocks: we instead pass both sets packed.
    // -> simpler: pass phase-0 and phase-1 dims, kernel picks via parity.

    // Pass R0,C0,M0 (phase 0); kernel derives phase-1 dims as R-? : we pass
    // per-phase through a small trick: relaunch is unnecessary because the
    // kernel receives Rp,Cp,Mp matching its phase from the arguments below.
    // We handle this by launching once with phase-0 dims for even z and
    // phase-1 dims for odd z using two arrays passed as scalars:
    // (kept simple: launch with max grid, kernel computes its own dims)

    // To keep a single launch, we pass both phases' geometry.
    // Rp/Cp/Mp arguments below correspond to phase (blockIdx.z & 1).

    fused_tconv_bias_tanh_kernel<<<grid, NTHREADS, 0, stream>>>(
        xc.data_ptr<float>(),
        wbuf.data_ptr<float>(),
        cc.data_ptr<float>(),
        sc.data_ptr<float>(),
        out.data_ptr<float>(),
        (int)batch, (int)IC, (int)H, (int)W,
        (int)OC, (int)OH, (int)OW,
        0 /*unused*/, 0 /*unused*/, 0 /*unused*/);

    return out;
}
"""

# NOTE: the kernel above needs per-phase geometry; we pass it explicitly by
# relaunching per phase inside the C++ (cleaner).  Redefine wrapper properly.

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

#define BM 64
#define BN 64
#define BK 16
#define NTHREADS 256

__global__ void build_wbuf_kernel(const float* __restrict__ w,
                                  float* __restrict__ wbuf,
                                  int IC, int OC) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = 2 * (IC * 4) * 2 * OC;
    if (idx >= total) return;
    int oc  = idx % OC;
    int qq  = (idx / OC) % 2;
    int kp  = (idx / (OC * 2)) % (IC * 4);
    int p   = idx / (OC * 2 * IC * 4);
    int ic = kp >> 2;
    int ta = (kp >> 1) & 1;
    int tb = kp & 1;
    int kh = 3 - p - 2 * ta;
    int kw = 3 - qq - 2 * tb;
    wbuf[((p * (IC * 4) + kp) * 2 + qq) * OC + oc] =
        w[((oc * IC + ic) * 4 + kh) * 4 + kw];
}

__global__ void __launch_bounds__(NTHREADS)
fused_tconv_bias_tanh_kernel(const float* __restrict__ x,
                             const float* __restrict__ wbuf,
                             const float* __restrict__ cbias,
                             const float* __restrict__ subbias,
                             float* __restrict__ out,
                             int batch, int IC, int H, int W,
                             int OC, int OH, int OW,
                             int R0, int C0, long M0,
                             int R1, int C1, long M1) {
    const int p   = blockIdx.z & 1;
    const int n   = blockIdx.z >> 1;
    const int Cp  = p ? C1 : C0;
    const long Mp = p ? M1 : M0;

    const int oc0 = blockIdx.y * BN;
    const long m0 = (long)blockIdx.x * BM;
    if (m0 >= Mp) return;

    __shared__ float As[2][2][BK][BM];
    __shared__ float Bs[2][2][BK][BN];

    const int tid = threadIdx.x;
    const int tm  = tid % 16;
    const int tn  = tid / 16;

    const int mi_ld   = tid % 64;
    const long m_ld   = m0 + mi_ld;
    const bool m_ok   = (m_ld < Mp);
    const long rc     = m_ok ? (m_ld / Cp) : 0;
    const long cc     = m_ok ? (m_ld - rc * Cp) : 0;
    const int rowbase = (int)rc + p - 1;
    const int colbase = (int)cc - 1;
    const long xbase  = ((long)n * IC) * H * W;

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

    const float* wbufp = wbuf + (long)p * (IC * 4) * 2 * OC;
    const int kchunks = (IC * 4) / BK;

    // ---- prologue: load chunk 0 ----
    {
        float av[8], bv[8];
        #pragma unroll
        for (int pass = 0; pass < 8; ++pass) {
            int e   = tid + pass * NTHREADS;
            int arr = e >> 10;
            int rem = e & 1023;
            int kp  = rem >> 6;
            {
                int ic  = kp >> 2;
                int ta  = (kp >> 1) & 1;
                int tb  = kp & 1;
                int row = rowbase + ta;
                int col = colbase + tb + arr;
                float v = 0.f;
                if (m_ok && row >= 0 && row < H && col >= 0 && col < W)
                    v = x[xbase + (long)ic * H * W + (long)row * W + col];
                av[pass] = v;
            }
            {
                int ocol = rem & 63;
                bv[pass] = wbufp[((long)kp * 2 + arr) * OC + oc0 + ocol];
            }
        }
        #pragma unroll
        for (int pass = 0; pass < 8; ++pass) {
            int e   = tid + pass * NTHREADS;
            int arr = e >> 10;
            int rem = e & 1023;
            int kp  = rem >> 6;
            int ocol = rem & 63;
            As[0][arr][kp][mi_ld] = av[pass];
            Bs[0][arr][kp][ocol]  = bv[pass];
        }
    }
    __syncthreads();

    for (int kc = 0; kc < kchunks; ++kc) {
        const int buf = kc & 1;
        float av[8], bv[8];
        const bool have_next = (kc + 1 < kchunks);
        if (have_next) {
            const int kn = kc + 1;
            #pragma unroll
            for (int pass = 0; pass < 8; ++pass) {
                int e   = tid + pass * NTHREADS;
                int arr = e >> 10;
                int rem = e & 1023;
                int kp  = rem >> 6;
                {
                    int kgl = kn * BK + kp;
                    int ic  = kgl >> 2;
                    int ta  = (kgl >> 1) & 1;
                    int tb  = kgl & 1;
                    int row = rowbase + ta;
                    int col = colbase + tb + arr;
                    float v = 0.f;
                    if (m_ok && row >= 0 && row < H && col >= 0 && col < W)
                        v = x[xbase + (long)ic * H * W + (long)row * W + col];
                    av[pass] = v;
                }
                {
                    int ocol = rem & 63;
                    int kgl  = kn * BK + kp;
                    bv[pass] = wbufp[((long)kgl * 2 + arr) * OC + oc0 + ocol];
                }
            }
        }
        #pragma unroll
        for (int kk = 0; kk < BK; ++kk) {
            float4 a0 = *reinterpret_cast<const float4*>(&As[buf][0][kk][tm * 4]);
            float4 a1 = *reinterpret_cast<const float4*>(&As[buf][1][kk][tm * 4]);
            float4 b0 = *reinterpret_cast<const float4*>(&Bs[buf][0][kk][tn * 4]);
            float4 b1 = *reinterpret_cast<const float4*>(&Bs[buf][1][kk][tn * 4]);
            #pragma unroll
            for (int i = 0; i < 4; ++i) {
                float av0 = (i == 0) ? a0.x : (i == 1) ? a0.y : (i == 2) ? a0.z : a0.w;
                float av1 = (i == 0) ? a1.x : (i == 1) ? a1.y : (i == 2) ? a1.z : a1.w;
                acc[0][i][0] = fmaf(av0, b0.x, acc[0][i][0]);
                acc[0][i][1] = fmaf(av0, b0.y, acc[0][i][1]);
                acc[0][i][2] = fmaf(av0, b0.z, acc[0][i][2]);
                acc[0][i][3] = fmaf(av0, b0.w, acc[0][i][3]);
                acc[1][i][0] = fmaf(av1, b1.x, acc[1][i][0]);
                acc[1][i][1] = fmaf(av1, b1.y, acc[1][i][1]);
                acc[1][i][2] = fmaf(av1, b1.z, acc[1][i][2]);
                acc[1][i][3] = fmaf(av1, b1.w, acc[1][i][3]);
            }
        }
        if (have_next) {
            __syncthreads();
            const int nbuf = buf ^ 1;
            #pragma unroll
            for (int pass = 0; pass < 8; ++pass) {
                int e   = tid + pass * NTHREADS;
                int arr = e >> 10;
                int rem = e & 1023;
                int kp  = rem >> 6;
                int ocol = rem & 63;
                As[nbuf][arr][kp][mi_ld] = av[pass];
                Bs[nbuf][arr][