```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# ---------------------------------------------------------------------------
# Math behind the fusion (kernel_size=3, stride=2, padding=1 transposed conv,
# followed by two stride-2 2x2x2 average pools):
#
#   convT output p receives x[i] with weight W[k] iff p = 2*i + k - 1.
#   Two avg pools with odd sizes (63 -> 31 -> 15) only use convT outputs
#   p in [0,59] per dim, i.e. final[q] = (1/64) * sum_{d in [0,3]^3} y[4q+d].
#   Gathering the 64 sub-terms gives a *valid* stride-2 kernel-3 convolution
#   on x[0..30] with an "effective" kernel E:
#       final[oc][q] = sum_ic sum_{j in [0,2]^3} E[oc][ic][j] * x[ic][2q+j]
#   where E folds W over the 64 pooled positions (and the 1/64 factor).
#   BatchNorm is per-channel affine, so it commutes with the average pool and
#   collapses into a per-channel scale/shift applied at the end.
#
# Train mode: batch statistics of the *full* convT output y (63^3 per slice)
# are computed on the fly (y is never materialized).  Each thread owns one py
# and an (even, odd) px pair, walks pz in (even, odd) pairs, register-blocks
# the <= 2x2x2x3 x values shared by the 4 y positions, and evaluates all 16
# output channels from shared-memory weights -- so each x tap is loaded once
# and reused for every tap/oc that needs it (~19x less load traffic than a
# direct evaluation).
# ---------------------------------------------------------------------------

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

#define ROWS 8
#define TOTAL_ROWS 225
#define X_TILE_FLOATS (ROWS * 3 * 3 * 3 * 31)
#define SLICE_ELEMS 250047   // 63*63*63

// Build effective kernel E[oc][ic][27] from conv-transpose W[ic][oc][27].
// Valid (d,j)->k mapping per dim: j=0 -> k={1,2}, j=1 -> k={0,1,2}, j=2 -> k={0}.
__global__ void build_eff_kernel(const float* __restrict__ W,
                                 float* __restrict__ E,
                                 int IC, int OC) {
    int oc = blockIdx.x / IC;
    int ic = blockIdx.x % IC;
    int j = threadIdx.x;              // 0..26
    if (j >= 27) return;
    int jz = j / 9, jy = (j / 3) % 3, jx = j % 3;

    int kzl[3], kyl[3], kxl[3], nz, ny, nx;
    if (jz == 0)      { kzl[0]=1; kzl[1]=2; nz=2; }
    else if (jz == 1) { kzl[0]=0; kzl[1]=1; kzl[2]=2; nz=3; }
    else              { kzl[0]=0; nz=1; }
    if (jy == 0)      { kyl[0]=1; kyl[1]=2; ny=2; }
    else if (jy == 1) { kyl[0]=0; kyl[1]=1; kyl[2]=2; ny=3; }
    else              { kyl[0]=0; ny=1; }
    if (jx == 0)      { kxl[0]=1; kxl[1]=2; nx=2; }
    else if (jx == 1) { kxl[0]=0; kxl[1]=1; kxl[2]=2; nx=3; }
    else              { kxl[0]=0; nx=1; }

    const float* wb = W + ((long)ic * OC + oc) * 27;
    float acc = 0.f;
    for (int a = 0; a < nz; ++a)
        for (int b = 0; b < ny; ++b)
            for (int c = 0; c < nx; ++c)
                acc += wb[kzl[a] * 9 + kyl[b] * 3 + kxl[c]];
    E[((long)oc * IC + ic) * 27 + j] = acc * (1.0f / 64.0f);
}

// Fully fused: effective stride-2 conv + BN affine + both avg pools.
// x: [N][IC][32][32][32], E: [OC][IC][27], out: [N][OC][15][15][15]
__global__ void fused_conv_bn_pool_kernel(
        const float* __restrict__ x,
        const float* __restrict__ E,
        const float* __restrict__ bn_w,
        const float* __restrict__ bn_b,
        const float* __restrict__ bn_mean,   // running_mean (eval) or batch mean (train)
        const float* __restrict__ bn_var,    // running_var  (eval) or batch biased var (train)
        const float* __restrict__ conv_bias,
        float* __restrict__ out,
        int N, int IC, int OC, float eps) {
    __shared__ float xs[X_TILE_FLOATS];
    __shared__ float Es[16 * 3 * 27];

    const int n   = blockIdx.z;
    const int row = blockIdx.x * ROWS + (threadIdx.x >> 4);  // 0..224
    const int qx  = threadIdx.x & 15;
    const int tid = threadIdx.x;

    // load effective kernel
    for (int i = tid; i < OC * IC * 27; i += blockDim.x)
        Es[i] = E[i];

    // load input tile: rows of this block -> x[n][ic][2qz+dz][2qy+dy][0..30]
    const int row0 = blockIdx.x * ROWS;
    const int nrows = min(ROWS, TOTAL_ROWS - row0);
    const long xbase = (long)n * IC * 32768;
    for (int idx = tid; idx < nrows * 3 * 3 * 3 * 31; idx += blockDim.x) {
        int xx  = idx % 31;
        int t   = idx / 31;
        int dy  = t % 3;
        int t2  = t / 3;
        int dz  = t2 % 3;
        int ic  = t2 / 3;
        int r   = ic / 3;  ic = ic % 3;
        int qz  = (row0 + r) / 15;
        int qy  = (row0 + r) % 15;
        int gz  = 2 * qz + dz;
        int gy  = 2 * qy + dy;
        xs[idx] = x[xbase + ((long)ic * 32 + gz) * 1024 + gy * 32 + xx];
    }
    __syncthreads();

    if (qx >= 15 || row >= TOTAL_ROWS) return;
    const int qz = row / 15;
    const int qy = row % 15;
    const int r  = threadIdx.x >> 4;

    for (int oc = 0; oc < OC; ++oc) {
        const float* ep = Es + oc * IC * 27;
        float acc = 0.f;
        for (int ic = 0; ic < IC; ++ic) {
            const float* xp = xs + ((r * 3 + ic) * 9) * 31 + 2 * qx;
            #pragma unroll
            for (int j = 0; j < 27; ++j) {
                int jz = j / 9, jy = (j / 3) % 3, jx = j % 3;
                acc += ep[j] * xp[(jz * 3 + jy) * 31 + jx];
            }
            ep += 27;
        }
        float scale = bn_w[oc] * rsqrtf(bn_var[oc] + eps);
        float shift = bn_b[oc] + scale * (conv_bias[oc] - bn_mean[oc]);
        out[(((long)n * OC + oc) * 225 + row) * 15 + qx] = acc * scale + shift;
    }
}

// Train-mode batch statistics of the full convT output y (63^3 per slice),
// computed on the fly: y is never written to global memory.
//   y[n][oc][pz][py][px] = bias[oc]
//       + sum_ic sum_{valid k taps} x[n][ic][iz][iy][ix] * W[ic][oc][k]
// Per dim: p even -> single tap (k=1, i=p/2);
//          p odd  -> two taps (k=0, i=(p+1)/2) and (k=2, i=(p-1)/2).
// block = (32, 8) threads: px = {2*tx, 2*tx+1}, py = 8*blockIdx.x + ty;
// pz walked in (even, odd) pairs so the <= 2x2x2 x values (x 3 input
// channels) sit in registers and feed all 4 y positions x 16 output channels.
// Per-block partials are reduced in shared memory and atomically accumulated.
__global__ void convt_stats_kernel(const float* __restrict__ x,
                                   const float* __restrict__ W,
                                   const float* __restrict__ bias,
                                   double* __restrict__ sum,
                                   double* __restrict__ sumsq,
                                   int IC, int OC) {
    __shared__ float Ws[16 * 3 * 27];
    __shared__ float bias_s[16];
    __shared__ double bs[256];
    __shared__ double bs2[256];

    const int n   = blockIdx.y;
    const int tid = threadIdx.x + (threadIdx.y << 5);

    for (int i = tid; i < OC * IC * 27; i += 256) {
        int oc = i / 81;
        int r  = i - oc * 81;
        int ic = r / 27;
        int k  = r - ic * 27;
        Ws[i] = W[((long)ic * OC + oc) * 27 + k];
    }
    if (tid < OC) bias_s[tid] = bias[tid];
    __syncthreads();

    const int tx = threadIdx.x;                     // 0..31
    const int py = (blockIdx.x << 3) + threadIdx.y; // 0..63
    const bool py_valid = (py <= 62);
    const bool px1_ok   = (tx <= 30);               // px = 2*tx+1 <= 62

    const int ix0 = tx;
    const int ix1 = (tx + 1 <= 31) ? (tx + 1) : 31;  // clamped; unused at tx=31
    const bool py_odd = (py & 1) != 0;
    const int iyA = py_odd ? ((py - 1) >> 1) : (py >> 1);  // ky=2 (odd) / ky=1 (even)
    const int iyB = py_odd ? ((py + 1) >> 1) : (py >> 1);  // ky=0 (odd); == iyA if even

    float acc_s[16], acc_q[16];
    #pragma unroll
    for (int c = 0; c < 16; ++c) { acc_s[c] = 0.f; acc_q[c] = 0.f; }

    const long xn = (long)n * IC * 32768;

    if (py_valid) {
        for (int t = 0; t < 32; ++t) {
            const bool tail = (t == 31);   // only even pz = 62
            // register-block the x values shared by the (up to) 4 y positions
            float xr[2][2][2][3];          // [dz][dy][dx][ic]
            {
                const int zlim = tail ? 0 : 1;
                #pragma unroll
                for (int dz = 0; dz < 2; ++dz) {
                    if (dz > zlim) break;
                    const long zoff = xn + (long)(t + dz) * 1024;
                    #pragma unroll
                    for (int dy = 0; dy < 2; ++dy) {
                        const long yoff = zoff + (long)(dy ? iyB : iyA) * 32;
                        #pragma unroll
                        for (int dx = 0; dx < 2; ++dx) {
                            const float* xp = x + yoff + (dx ? ix1 : ix0);
                            #pragma unroll
                            for (int ic = 0; ic < 3; ++ic)
                                xr[dz][dy][dx][ic] = xp[(long)ic * 32768];
                        }
                    }
                }
            }

            float y0[16], y1[16], y2[16], y3[16];
            #pragma unroll
            for (int c = 0; c < 16; ++c) {
                float b = bias_s[c];
                y0[c] = b; y1[c] = b; y2[c] = b; y3[c] = b;
            }

            #pragma unroll
            for (int oc = 0; oc < 16; ++oc) {
                const float* wp = Ws + oc * 81;
                #pragma unroll
                for (int ic = 0; ic < 3; ++ic) {
                    const float* w = wp + ic * 27;
                    if (py_odd) {
                        float xa0 = xr[0][0][0][ic], xa1 = xr[0][0][1][ic];
                        float xb0 = xr[0][1][0][ic], xb1 = xr[0][1][1][ic];
                        // even pz (kz=1)
                        y0[oc] += w[16] * xa0 + w[7] * xb0;              // ky=2,ky=0 (kx=1)
                        y1[oc] += w[17] * xa0 + w[15] * xa1              // ky=2
                                + w[8]  * xb0 + w[6]  * xb1;             // ky=0
                        if (!tail) {
                            float xc0 = xr[1][0][0][ic], xc1 = xr[1][0][1][ic];
                            float xd0 = xr[1][1][0][ic], xd1 = xr[1][1][1][ic];
                            // odd pz: kz=2 (dz=0), kz=0 (dz=1)
                            y2[oc] += w[25] * xa0 + w[16] * xb0
                                    + w[7]  * xc0 + w[7 - 7 + 1] * 0.f;  // placeholder cleaned below
                            y2[oc] += w[1] * xc0 + w[1 + 9] * 0.f;       // cleaned below
                            y2[oc] += w[16 + 9] * 0.f;
                            y3[oc] += 0.f;
                        }
                    } else {
                        float xa0 = xr[0][0][0][ic], xa1 = xr[0][0][1][ic];
                        y0[oc] += w[13] * xa0;                           // ky=1, kx=1
                        y1[oc] += w[14] * xa0 + w[12] * xa1;             // kx=2, kx=0
                    }
                }
            }
        }
    }

    // placeholder (never compiled -- replaced below)
    bs[tid] = 0.0; bs2[tid] = 0.0;
    __syncthreads();
}

// compute batch mean/var and update running stats like aten batch_norm
__global__ void bn_finalize_kernel(const double* __restrict__ sum,
                                   const double* __restrict__ sumsq,
                                   float* __restrict__ running_mean,
                                   float* __restrict__ running_var,
                                   float* __restrict__ out_mean,
                                   float* __restrict__ out_var,
                                   long count, double momentum, float eps, int C) {
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= C) return;
    double m  = sum[c] / (double)count;
    double sq = sumsq[c] / (double)count;
    double v  = sq - m * m;
    if (v < 0.0) v = 0.0;
    out_mean[c] = (float)m;
    out_var[c]  = (float)v;
    double unb = (count > 1) ? v * ((double)count / (double)(count - 1)) : 0.0;
    running_mean[c] = (float)((1.0 - momentum) * (double)running_mean[c] + momentum * m);
    running_var[c]  = (float)((1.0 - momentum) * (double)running_var[c]  + momentum * unb);
}

torch::Tensor build_eff(torch::Tensor W, int64_t IC, int64_t OC) {
    auto E = torch::empty({OC, IC, 27}, W.options());
    auto stream = at::cuda::getCurrentCUDAStream();
    dim3 grid(OC * IC);
    build_eff_kernel<<<grid, 32, 0, stream>>>(W.data_ptr<float>(), E.data_ptr<float>(), (int)IC, (int)OC);
    return E;
}

torch::Tensor fused_eval(torch::Tensor x, torch::Tensor E,
                         torch::Tensor bn_w, torch::Tensor bn_b,
                         torch::Tensor bn_mean, torch::Tensor bn_var,
                         torch::Tensor conv_bias,
                         int64_t N, int64_t IC, int64_t OC, double eps) {
    auto out = torch::empty({N, OC, 15, 15, 15}, x.options());
    auto stream = at::cuda::getCurrentCUDAStream();
    dim3 grid((TOTAL_ROWS + ROWS - 1) / ROWS, 1, N);
    fused_conv_bn_pool_kernel<<<grid, 128, 0, stream>>>(
        x.data_ptr<float>(), E.data_ptr<float>(),
        bn_w.data_ptr<float>(), bn_b.data_ptr<float>(),
        bn_mean.data_ptr<float>(), bn_var.data_ptr<float>(),
        conv_bias.data_ptr<float>(), out.data_ptr<float>(),
        (int)N, (int)IC, (int)OC, (float)eps);
    return out;
}

// train-mode: on-the-fly convT batch stats (y never materialized),
// then fused effective-conv + BN + pools with the batch stats.
torch::Tensor fused_train(torch::Tensor x, torch::Tensor E, torch::Tensor W,
                          torch::Tensor bn_w, torch::Tensor bn_b,
                          torch::Tensor running_mean, torch::Tensor running_var,
                          torch::Tensor conv_bias,
                          int64_t N, int64_t IC, int64_t OC,
                          double eps, double momentum) {
    auto stream = at::cuda::getCurrentCUDAStream();
    auto opts_d = torch::TensorOptions().dtype(torch::kFloat64).device(x.device());
    auto sum    = torch::zeros({OC}, opts_d);
    auto sumsq  = torch::zeros({OC}, opts_d);

    dim3 sgrid(8, (unsigned)N, 1);
    dim3 sblock(32, 8);
    convt_stats_kernel<<<sgrid, sblock, 0, stream>>>(
        x.data_ptr<float>(), W.data_ptr<float>(), conv_bias.data_ptr<float>(),
        sum.data_ptr<double>(), sumsq.data_ptr<double>(), (int)IC, (int)OC);

    auto bmean = torch::empty({OC}, x.options());
    auto bvar  = torch::empty({OC}, x.options());
    long count = (long)N * SLICE_ELEMS;
    bn_finalize_kernel<<<(unsigned)((OC + 63) / 64), 64, 0, stream>>>(
        sum.data_ptr<double>(), sumsq.data_ptr<double>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        bmean.data_ptr<float>(), bvar.data_ptr<float>(),
        count, momentum, (float)eps, (int)OC);

    auto out = torch::empty({N, OC, 15, 15, 15}, x.options());
    dim3 grid((TOTAL_ROWS + ROWS - 1) / ROWS, 1, N);
    fused_conv_bn_pool_kernel<<<grid, 128, 0, stream>>>(
        x.data_ptr<float>(), E.data_ptr<float>(),
        bn_w.data_ptr<float>(), bn_b.data_ptr<float>(),
        bmean.data_ptr<float>(), bvar.data_ptr<float>(),
        conv_bias.data_ptr<float>(), out.data_ptr<float>(),
        (int)N, (int)IC, (int)OC, (float)eps);
    return out;
}
"""
```

Wait — I started writing placeholder garbage in the kernel. I need to write the FMA section cleanly and completely. Let me rewrite the whole thing properly in my head and then produce the final file.

Let me carefully define, per (ic, oc), for py_odd, one pz-pair t:

y0 = y(2t, py, 2tx): pz even → kz=1, iz=t; py odd; px even → kx=1, ix=tx.
taps: kz=1 (iz=t), ky∈{2 (iy=iyA), 0 (iy=iyB)}, kx=1 (ix=tx).
y0 = Σ_ic w[1*9 + 2*3 + 1]·xr[0][0][0] + w[1*9 + 0*3 + 1]·xr[0][1][0]
   = w[16]·xr[0][0][0] + w[7]·xr[0][1][0]. ✓

y1 = y(2t, py, 2tx+1): pz even kz=1 (iz=t), py odd, px odd → kx∈{2 (ix=tx), 0 (ix=tx+1)}.
y1 = Σ w[1*9+2*3+2]·xr[0][0][0] + w[1*9+2*3+0]·xr[0][0][1] + w[1*9+0*3+2]·xr[0][1][0] + w[1*9+0*3+0]·xr[0][1][1]
   = w[17]·xr[0][0][0] + w[15]·xr[0][0][1] + w[8]·xr[0][1][0] + w[6]·xr[0][1][1]. ✓

y2 = y(2t+1, py, 2tx): pz odd → kz∈{2 (iz=t), 0 (iz=t+1)}, py odd, px even kx=1.
y2 = w[2*9+2*3+1]·xr[0][0][0] + w[2*9+0*3+1]·xr[0][1][0] + w[0*9+2*3+1]·xr[1][0][0] + w[0*9+0*3+1]·xr[1][1][0]
   = w[25]·xr[0][0][0] + w[19]·xr[0][1][0] + w[7]·xr[1][0][0] + w[1]·xr[1][1][0]. ✓

y3 = y(2t+1, py, 2tx+1): pz odd, py odd, px odd → 8 taps:
y3 = w[2*9+2*3+2]·xr[0][0][0] + w[2*9+2*3+0]·xr[0][0][1]
   + w[2*9+0*3+2]·xr[0][1][0] + w[2*9+0*3+0]·xr[0][1][1]
   + w[0*9+2*3+2]·xr[1][0][0] + w[0*9+2*3+0]·xr[1][0][1]
   + w[0*9+0*3+2]·xr[1][1][0] + w[0*9+0*3+0]·xr[1][1][1]
   = w[26]·xr000 + w[24]·xr001 + w[20]·xr010 + w[18]·xr011
   + w[8]·xr100  + w[6]·xr101  + w[2]·xr110  + w[0]·xr111. ✓

py even:
y0 = w[1*9+1*3+1]·xr[0][0][0] = w[13]·xr000.
y1 = w[14]·xr000 + w[12]·xr001.
y2 = w[2*9+1*3+1]·xr000 + w[0*9+1*3+1]·xr100 = w[22]·xr000 + w[4]·xr100.
y3 = w[2*9+1*3+2]·xr000 + w[2*9+1*3+0]·xr001 + w[0*9+1*3+2]·xr100 + w[0*9+1*3+0]·xr101
   = w[23]·xr000 + w[21]·xr001 + w[5]·xr100 + w[3]·xr101. ✓

Accumulation per pair: acc_s[oc] += y0+y1 (+y2+y3 if !tail); acc_q[oc] += y0²+y1² (+y2²+y3² if !tail); if tx==31: only y0,y2 (px even); i.e., skip y1,y3.

Guards: px1_ok = tx ≤ 30; tail = t == 31 (skip y2,y3 since pz=2t+1=63 invalid).

For tx=31: xr[*][*][1][*] loads are clamped to ix=31 (same as ix0) — wasted but harmless, values unused (we skip y1/y3 contributions). But wait — at tx=31 we still LOAD xr[..][1] — fine.

For tail t=31: we skip loading dz=1 (zlim=0). The `break` inside `#pragma unroll` on runtime condition — that will compile (unrolled with branches). Alternatively use `if (dz <= zlim)`. Cleaner: `for dz... if (tail && dz==1) continue;` Hmm with #pragma unroll and runtime tail — produces both paths; fine. Actually simpler: always load dz=1 with clamped index iz = min(t+1, 31): then no branch; loads wasted at tail but harmless: xr[1] loads iz=min(t+1,31)=31 for t=31 → reads iz=31 plane (in bounds ✓). Let's do: `const int iz1 = (t + 1 <= 31) ? t + 1 : 31;` and always load. Simpler control flow, static unroll.

Now also py=63 invalid: we skip the whole loop — but those threads still need to participate in __syncthreads at reduction (the load sync at top already done before branching ✓; reduction syncs after — must be OUTSIDE the py_valid guard ✓).

W load loop: OC*IC*27 = 1296, with hardcoded 16/3? Use passed IC/OC but they are 3/16 guaranteed by Python guard. Use runtime IC/OC with shared arrays sized 16*3*27 — fine since guard ensures. In kernel, loops `for i < OC*IC*27` ✓; `Ws[i] = W[(ic*OC+oc)*27+k]` ✓. bias_s sized 16, `if (tid < OC)` ✓. acc arrays sized 16, loop `#pragma unroll for c < 16` — use OC? For unroll benefits use compile-time 16; guarded OC==16 ✓. I'll hardcode 16/3 in inner loops with constants but keep parameters for outer sizes. Since Python guarantees IC=3, OC=16, hardcoding is safe.

Precision of accumulators: per thread 126 y values → float acc fine.

Block reduction: like v2: bs/bs2 256 doubles, loop oc: write acc_s[oc] etc. Actually per oc do: bs[tid] = acc_s[oc]; bs2[tid] = acc_q[oc]; __syncthreads(); tree 128..1; __syncthreads(); tid==0 atomicAdd. That's v2's pattern in a loop over 16 oc: 16 × (2 syncthreads × 8 levels...) — v2 did tree with syncs each level (8 levels × 16 oc = 128 syncs + …). It's fine (reduction cost ~ 16 × 256×8B writes... trivial).

Optimization: convert acc to double once: bs[tid] = (double)acc_s[oc]. ✓

Now the launch: grid (8, N, 1), block (32, 8). 512 blocks for N=64. Each block: 256 threads × ~31 pairs × (loads ~12-24 + FMAs 432-864... wait per thread per pair: FMAs = 16 oc × (py odd: 3+6+6+12... let me recount FMAs per oc per pair: py odd: y0:2... hold on — per ic: y0: 2 taps, y1: 4, y2: 4, y3: 8 → 18 FMAs per ic per oc → ×3 ic = 54 FMAs per oc per pair → ×16 oc = 864 FMAs per pair (py odd), py even: per ic: 1+2+2+4 = 9 → 27 per oc → 432 per pair. Plus square/sum accumulation: 16 oc × 4-6 FMA-ish ops ≈ 80-96. So ~500-950 FMAs per pair per thread → total = 512 blocks × 256 thr × 31 × avg(432+864)/2×... per n: threads = 8×256 = 2048 → per n per pair: 2048×avg648 ≈ 1.33M × 31 pairs... ≈ 41M × 64 n = 2.6G ✓ checks out.

Registers: xr 24 + y0..y3 64 + acc 32 + misc ≈ 130+ regs — occupancy 256 thr × 130 = 33K regs → 1 block/SM per reg file 64K? 65536/33280 ≈ 1.9 → 1 block... hmm that's low: 256 threads/SM only?! That's bad — only 108×256 = 27K threads resident; latency hiding poor. Hmm. Wait: 64K registers per SM (A100: 65536 per SM... actually 64K 32-bit registers per SM). 130 regs × 256 threads = 33,280 → 1 block/SM (2 would need 66,560 > 65,536). So 256 threads/SM = 8 warps — very low occupancy. But ILP within thread is huge (16 independent oc chains × 4 y). FMA dual... A100 FP32 64/cyc/SM; with 8 warps and high ILP, can sustain? Each warp needs to issue 2 FMA/cyc total across 8 warps to hit 64/32=2 per warp-scheduler... 4 schedulers, 2 warps each → each scheduler needs 1 FMA/cyc from 2 warps — with ILP 16, achievable if no stalls. Loads: 12-24 per pair with L1 hits mostly... risky but plausible ~60-80% peak. Alternatively reduce y registers: accumulate into acc_s/acc_q immediately per oc instead of keeping y0..y3 arrays for all 16 oc: restructure loop:

for oc: compute y0,y1,y2,y3 (4 regs), then acc_s[oc] += ..., acc_q[oc] += ... → y regs = 4 instead of 64! Registers: xr 24 + acc 32 + 4 + misc ≈ 70-80 regs → 80×256 = 20,480 → 3 blocks/SM = 768 threads ✓. ILP: within oc iteration, 4 y chains + xr reuse — acc_s[oc] dependency serializes across pairs but 16 oc × ... per pair: oc loop gives 16 independent accumulator pairs → ILP fine.

But xr must be loaded before oc loop ✓ (24 loads per pair). Then oc loop reads Ws from shared: per oc per ic: 18 shared loads (py odd) → 54 per oc → 864 shared loads per pair per thread?! Shared throughput: 32 floats/cyc/SM → total shared load instr: 512 blocks×256×31×864/32(warp)... = per warp per pair 864 shared-load instr... total = 512×8 warps×31×864 ≈ 110M warp shared-loads / 108 SMs ≈ 1M per SM ≈ 1M cycles?!?! No wait — shared load per warp instruction delivers 32 values (one per lane) — the count 864 is per THREAD scalar loads = per warp 864 warp-instructions (each 32 lanes × 4B = 128B). Shared bandwidth per SM: 128B/cyc = one warp-instruction per cycle. So per SM: total warp shared instrs = (512×8×31×864)/108 ≈ 1.016M cycles ≈ 720µs!!! WAY too slow. Shared-bound disaster.

Fix: W into registers? 54 floats per oc — can't (indexed by oc loop... could unroll oc fully → 16×81 = 1296 floats — no).

Alternative: swap roles — keep x in shared, W in registers! Per thread: W values needed = those for its (parity pattern) taps: py odd: per oc per ic: 18 W values (taps) — with oc loop 16 → 18×16 = 288 per ic... no: per ic the tap SET is fixed (18 weights per (ic,oc)) → total W per thread = 3 ic × 