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


# ============================================================================
# Custom CUDA kernels for: ConvTranspose3d(3->16, k=3, s=2, p=1) + BatchNorm3d
# + AvgPool3d(2) + AvgPool3d(2).
#
# Key algorithmic improvements over the previous version:
#  * The two 2x average pools exactly cancel the transposed-conv upsampling:
#    the pooled (pre-BN) output is a plain stride-2 3D cross-correlation of
#    the input with a combined 3x3x3 kernel K built from the conv weights,
#        K[i] = [w1+w2, w0+w1+w2, w0]  (outer product over the 3 axes).
#    So the 15^3 output is computed directly from the 32^3 input (~10x fewer
#    FLOPs than materializing the 63^3 volume).
#  * Training mode needs BatchNorm statistics over the full 63^3 conv output.
#    The per-channel sum is linear in the input and collapses to 27 position
#    classes of input-voxel sums.  The sum of squares is computed exactly via
#    input self-correlations (offsets in {-1,0,1}^3) plus an
#    inclusion-exclusion boundary correction (faces/edges/corners) that
#    accounts for the padding=1 crop.  This avoids ever forming the 63^3
#    volume (~5x fewer FLOPs than the direct computation).
#  * Inference mode is one tiny weight-build kernel + one fused conv kernel.
# ============================================================================

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

#define CO 16
#define CI 3
#define IDIM 32

// ---------------------------------------------------------------------------
// buildW: derive all combined weights from the conv-transpose weight
// w[ic][oc][kd][kh][kw] (layout [CI][CO][3][3][3]).
// ---------------------------------------------------------------------------
__global__ void buildW_kernel(
    const float* __restrict__ w,
    float* __restrict__ K,    // [27][CI][CO]   pool conv weights
    float* __restrict__ WT,   // [27][CI][CO]   per-position-class tap sums
    float* __restrict__ M3,   // [27][CI][CI][CO] volume correlation weights
    float* __restrict__ M2,   // [6][9][CI][CI][CO] face correlation weights
    float* __restrict__ M1,   // [12][3][CI][CI][CO] edge correlation weights
    float* __restrict__ M0)   // [8][CI][CI][CO] corner weights
{
    const int tid = threadIdx.x;
    const int nd = blockIdx.x;   // 0..2 axis of face / edge free axis
    const int s  = blockIdx.y;   // side 0/1
    const int kd_fix = s * 2;

    __shared__ float Nm[9];
    __shared__ float Kc[27][CI][CO];
    __shared__ float WTc[27][CI][CO];
    __shared__ float M3c[27][CI][CI][CO];
    __shared__ float M2c[9][CI][CI][CO];
    __shared__ float M1c[3][CI][CI][CO];
    __shared__ float M0c[CI][CI][CO];

    if (tid < 9) {
        const int i = tid / 3, k = tid % 3;
        Nm[tid] = (float)(((i == 0 && k >= 1) || (i == 1) || (i == 2 && k == 0)) ? 1 : 0);
    }
    __syncthreads();

    if (nd == 0) {
        // ---- combined pool kernel K + position-class sums WT + M3 ----
        for (int idx = tid; idx < 27 * CI * CO; idx += blockDim.x) {
            const int i  = idx / (CI * CO);
            const int r  = idx % (CI * CO);
            const int ic = r / CO;
            const int oc = r % CO;
            const int a = i / 9, b = (i / 3) % 3, c = i % 3;
            float acc = 0.f;
            #pragma unroll
            for (int kx = 0; kx < 3; ++kx)
                #pragma unroll
                for (int ky = 0; ky < 3; ++ky)
                    #pragma unroll
                    for (int kz = 0; kz < 3; ++kz)
                        acc += Nm[a * 3 + kx] * Nm[b * 3 + ky] * Nm[c * 3 + kz]
                             * w[((ic * CO + oc) * 3 + kx) * 9 + ky * 3 + kz];
            Kc[i][ic][oc] = acc;
            K[idx] = acc;
        }
        for (int idx = tid; idx < 27 * CI * CO; idx += blockDim.x) {
            const int i  = idx / (CI * CO);
            const int r  = idx % (CI * CO);
            const int ic = r / CO;
            const int oc = r % CO;
            const int a = i / 9, b = (i / 3) % 3, c = i % 3;
            const int sd = (a == 0) ? 1 : (a == 2 ? 0 : 2);
            const int sh = (b == 0) ? 1 : (b == 2 ? 0 : 2);
            const int sw = (c == 0) ? 1 : (c == 2 ? 0 : 2);
            float acc = 0.f;
            for (int kx = sd; kx < sd + 3; ++kx)
                for (int ky = sh; ky < sh + 3; ++ky)
                    for (int kz = sw; kz < sw + 3; ++kz)
                        acc += w[((ic * CO + oc) * 3 + (kx % 3)) * 9 + (ky % 3) * 3 + (kz % 3)];
            WTc[i][ic][oc] = acc;
            WT[idx] = acc;
        }
        __syncthreads();
        for (int idx = tid; idx < 27 * CI * CI * CO; idx += blockDim.x) {
            const int off = idx / (CI * CI * CO);
            int r = idx % (CI * CI * CO);
            const int ic = r / (CI * CO);
            r %= (CI * CO);
            const int jc = r / CO;
            const int oc = r % CO;
            const int od = off / 9 - 1, oh = (off / 3) % 3 - 1, ow = off % 3 - 1;
            float acc = 0.f;
            for (int i = 0; i < 27; ++i) {
                const float k1 = Kc[i][ic][oc];
                const int a = i / 9, b = (i / 3) % 3, c = i % 3;
                const int a2 = a + od, b2 = b + oh, c2 = c + ow;
                if ((unsigned)a2 < 3u && (unsigned)b2 < 3u && (unsigned)c2 < 3u)
                    acc += k1 * Kc[a2 * 9 + b2 * 3 + c2][jc][oc];
            }
            M3c[off][ic][jc][oc] = acc;
            M3[idx] = acc;
        }
        __syncthreads();
    }

    if (nd < 3) {
        // ---- face weights M2 (3 orientations x 2 sides) ----
        for (int idx = tid; idx < 9 * CI * CI * CO; idx += blockDim.x) {
            const int off = idx / (CI * CI * CO);
            int r = idx % (CI * CI * CO);
            const int ic = r / (CI * CO);
            r %= (CI * CO);
            const int jc = r / CO;
            const int oc = r % CO;
            const int o1 = off / 3 - 1, o2 = off % 3 - 1;
            float acc = 0.f;
            #pragma unroll
            for (int u = 0; u < 3; ++u) {
                #pragma unroll
                for (int v = 0; v < 3; ++v) {
                    const int u2 = u + o1, v2 = v + o2;
                    if ((unsigned)u2 < 3u && (unsigned)v2 < 3u) {
                        float w1, w2;
                        if (nd == 0) {
                            w1 = w[((ic * CO + oc) * 3 + kd_fix) * 9 + u * 3 + v];
                            w2 = w[((jc * CO + oc) * 3 + kd_fix) * 9 + u2 * 3 + v2];
                        } else if (nd == 1) {
                            w1 = w[((ic * CO + oc) * 3 + u) * 9 + kd_fix * 3 + v];
                            w2 = w[((jc * CO + oc) * 3 + u2) * 9 + kd_fix * 3 + v2];
                        } else {
                            w1 = w[((ic * CO + oc) * 3 + u) * 9 + v * 3 + kd_fix];
                            w2 = w[((jc * CO + oc) * 3 + u2) * 9 + v2 * 3 + kd_fix];
                        }
                        acc += w1 * w2;
                    }
                }
            }
            if (nd == 0) M2c[off][ic][jc][oc] = acc;
            M2[(((nd * 2 + s) * 9) + off) * CI * CI * CO + idx] = acc;
        }
        __syncthreads();
    }

    // ---- edge weights M1 (blockIdx.x = free axis + 3) ----
    if (nd >= 3) {
        const int free_ax = nd - 3;
        const int sh2 = blockIdx.y;              // 0..3 fixed sides
        const int s1 = sh2 >> 1, s2 = sh2 & 1;
        const int kf1 = s1 * 2, kf2 = s2 * 2;
        for (int idx = tid; idx < 3 * CI * CI * CO; idx += blockDim.x) {
            const int off = idx / (CI * CI * CO);
            int r = idx % (CI * CI * CO);
            const int ic = r / (CI * CO);
            r %= (CI * CO);
            const int jc = r / CO;
            const int oc = r % CO;
            const int o = off - 1;
            float acc = 0.f;
            #pragma unroll
            for (int u = 0; u < 3; ++u) {
                const int u2 = u + o;
                if ((unsigned)u2 < 3u) {
                    float w1, w2;
                    if (free_ax == 0) {
                        w1 = w[((ic * CO + oc) * 3 + u) * 9 + kf1 * 3 + kf2];
                        w2 = w[((jc * CO + oc) * 3 + u2) * 9 + kf1 * 3 + kf2];
                    } else if (free_ax == 1) {
                        w1 = w[((ic * CO + oc) * 3 + kf1) * 9 + u * 3 + kf2];
                        w2 = w[((jc * CO + oc) * 3 + kf1) * 9 + u2 * 3 + kf2];
                    } else {
                        w1 = w[((ic * CO + oc) * 3 + kf1) * 9 + kf2 * 3 + u];
                        w2 = w[((jc * CO + oc) * 3 + kf1) * 9 + kf2 * 3 + u2];
                    }
                    acc += w1 * w2;
                }
            }
            M1[(((free_ax * 4 + sh2) * 3) + off) * CI * CI * CO + idx] = acc;
        }
        // ---- corner weights M0 (only meaningful from block (3,0)) ----
        if (free_ax == 0 && sh2 == 0) {
            for (int idx = tid; idx < 8 * CI * CI * CO; idx += blockDim.x) {
                const int cn = idx / (CI * CI * CO);
                int r = idx % (CI * CI * CO);
                const int ic = r / (CI * CO);
                r %= (CI * CO);
                const int jc = r / CO;
                const int oc = r % CO;
                const int kdx = (cn >> 2) & 1, khy = (cn >> 1) & 1, kwz = cn & 1;
                const float w1 = w[((ic * CO + oc) * 3 + kdx * 2) * 9 + khy * 2 * 3 + kwz * 2];
                const float w2 = w[((jc * CO + oc) * 3 + kdx * 2) * 9 + khy * 2 * 3 + kwz * 2];
                M0c[ic][jc][oc] = w1 * w2;
                M0[idx] = w1 * w2;
            }
        }
    }

    // ---- scaled pool-conv output weights: K * gamma*invstd/64, from block 0
    if (blockIdx.x == 0 && blockIdx.y == 1) {
        const int oc = tid;
        if (oc < CO) {
            // placeholder; real per-forward scale happens in the wrappers
        }
    }
}

// ---------------------------------------------------------------------------
// poolConv: out[n][oc][bd][bh][bw] = sum_ic sum_{27} K * x[n][ic][2bd+i,..]
// K layout: [27][CI][CO]; optionally applies scale/shift (BN folded).
// Input patch per block: 10^3 voxels x 3 channels in shared memory.
// ---------------------------------------------------------------------------
__global__ void poolConv_kernel(
    const float* __restrict__ x,
    const float* __restrict__ K,
    const float* __restrict__ scale,   // [CO] or nullptr
    const float* __restrict__ shift,   // [CO] or nullptr
    float* __restrict__ out,
    int nimg, int od)
{
    __shared__ float sx[10 * 10 * 10 * CI];
    __shared__ float sk[27 * CI * CO];
    __shared__ float ssc[CO];
    __shared__ float ssh[CO];

    const int tid = threadIdx.x;   // 125 threads
    const int n = blockIdx.x;
    const int b0 = blockIdx.y * 5;
    const int h0 = blockIdx.z * 5;

    for (int i = tid; i < 27 * CI * CO; i += 125) sk[i] = K[i];
    if (tid < CO) {
        ssc[tid] = scale ? scale[tid] : 1.0f;
        ssh[tid] = shift ? shift[tid] : 0.0f;
    }

    const float* xb = x + (long)n * CI * IDIM * IDIM * IDIM;
    const bool interior = (b0 > 0) && (b0 + 5 < od) && (h0 > 0) && (h0 + 5 < od);
    if (interior) {
        for (int i = tid; i < 1000 * CI; i += 125) {
            const int ic = i / 1000;
            const int a = (i % 1000) / 100;
            const int b = (i % 100) / 10;
            const int c = i % 10;
            sx[i] = xb[((ic * IDIM + (2 * b0 - 2 + a)) * IDIM + (2 * h0 - 2 + b)) * IDIM + c];
        }
    } else {
        for (int i = tid; i < 1000 * CI; i += 125) {
            const int ic = i / 1000;
            const int a = (i % 1000) / 100;
            const int b = (i % 100) / 10;
            const int c = i % 10;
            const int gi = 2 * b0 - 2 + a, gj = 2 * h0 - 2 + b, gk = c;
            float v = 0.f;
            if ((unsigned)gi < 32u && (unsigned)gj < 32u && gk < 32)
                v = xb[((ic * IDIM + gi) * IDIM + gj) * IDIM + gk];
            sx[i] = v;
        }
    }
    __syncthreads();

    const int lz = tid % 5, ly = (tid / 5) % 5, lx = tid / 25;
    const int bd = b0 + lx, bh = h0 + ly, bw = lz;
    if (bd >= od || bh >= od || bw >= od) return;

    const int base = lx * 200 + ly * 20 + lz * 2;
    float xv[27];
    #pragma unroll
    for (int i = 0; i < 3; ++i)
        #pragma unroll
        for (int j = 0; j < 3; ++j)
            #pragma unroll
            for (int k = 0; k < 3; ++k)
                xv[i * 9 + j * 3 + k] = sx[base + i * 100 + j * 10 + k];

    float acc[CO];
    #pragma unroll
    for (int oc = 0; oc < CO; ++oc) acc[oc] = 0.f;

    #pragma unroll
    for (int t = 0; t < 27; ++t) {
        const float* kp = sk + t * CI * CO;
        const float x0 = xv[t], x1 = xv[t + 1000], x2 = xv[t + 2000];
        #pragma unroll
        for (int oc = 0; oc < CO; ++oc)
            acc[oc] += x0 * kp[oc] + x1 * kp[CO + oc] + x2 * kp[2 * CO + oc];
    }

    const long oidx = (((long)n * CO) * od + bd) * od * od + (long)bh * od + bw;
    #pragma unroll
    for (int oc = 0; oc < CO; ++oc)
        out[oidx * CO + oc] = acc[oc] * ssc[oc] + ssh[oc];
}

// channel-last -> NCDHW for the pool output
__global__ void cl_to_nc(const float* __restrict__ in, float* __restrict__ out,
                         long total, int osize)
{
    const long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total) return;
    const int oc = i % CO;
    const long v = i / CO;
    out[(v / osize) * CO * osize + (long)oc * osize + (v % osize)] = in[i];
}

// NCDHW -> channel-last
__global__ void nc_to_cl(const float* __restrict__ in, float* __restrict__ out,
                         long total, int osize)
{
    const long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total) return;
    const int oc = i % CO;
    const long v = i / CO;
    out[i] = in[(v / osize) * CO * osize + (long)oc * osize + (v % osize)];
}

// fused channel-last BN scale/shift for the raw pooled sums
__global__ void bn_scale_cl(const float* __restrict__ in, float* __restrict__ out,
                            const float* __restrict__ scale,
                            const float* __restrict__ shift, long total)
{
    const long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total) return;
    const int oc = i % CO;
    out[i] = in[i] * scale[oc] + shift[oc];
}

// ---------------------------------------------------------------------------
// xsClasses: XS[ic][27 position classes] summed over batch (double accum)
// ---------------------------------------------------------------------------
__global__ void xsClasses_kernel(const float* __restrict__ x, long total, double* __restrict__ XS)
{
    __shared__ double sxs[81];
    for (int i = threadIdx.x; i < 81; i += blockDim.x) sxs[i] = 0.0;
    __syncthreads();

    const long stride = (long)gridDim.x * blockDim.x;
    for (long idx = (long)blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += stride) {
        const int w = idx % IDIM;
        const int h = (idx / IDIM) % IDIM;
        const int d = (idx / (IDIM * IDIM)) % IDIM;
        const int ic = (idx / (IDIM * IDIM * IDIM)) % CI;
        const int cd = (d == 0) ? 0 : (d == 31 ? 2 : 1);
        const int ch = (h == 0) ? 0 : (h == 31 ? 2 : 1);
        const int cw = (w == 0) ? 0 : (w == 31 ? 2 : 1);
        atomicAdd(&sxs[ic * 27 + cd * 9 + ch * 3 + cw], (double)x[idx]);
    }
    __syncthreads();
    for (int i = threadIdx.x; i < 81; i += blockDim.x) {
        if (sxs[i] != 0.0) atomicAdd(&XS[i], sxs[i]);
    }
}

// ---------------------------------------------------------------------------
// corr3: C3[off][ic][jc] = sum over batch, p, p+off in range of x*x (double)
// grid: (chunks, 27 offsets); block 128 threads.
// ---------------------------------------------------------------------------
__global__ void corr3_kernel(const float* __restrict__ x, int nimg,
                             int chunks_per_img, double* __restrict__ C3)
{
    const int off = blockIdx.y;
    const int od = off / 9 - 1, oh = (off / 3) % 3 - 1, ow = off % 3 - 1;

    float acc[CI * CI];
    #pragma unroll
    for (int i = 0; i < CI * CI; ++i) acc[i] = 0.f;

    const long per_img = (long)CI * IDIM * IDIM * IDIM;
    const long chunk = (per_img + chunks_per_img - 1) / chunks_per_img;

    for (int n = 0; n < nimg; ++n) {
        const long base = (long)n * per_img + (long)blockIdx.x * chunk;
        const long end = min(base + chunk, (long)(n + 1) * per_img);
        for (long idx = base + threadIdx.x; idx < end; idx += blockDim.x) {
            const long loc = idx % per_img;
            const int w = loc % IDIM;
            const int h = (loc / IDIM) % IDIM;
            const int d = (loc / (IDIM * IDIM)) % IDIM;
            const int w2 = w + ow, h2 = h + oh, d2 = d + od;
            if ((unsigned)w2 < 32u && (unsigned)h2 < 32u && (unsigned)d2 < 32u) {
                const long sh2 = ((long)(d2 * IDIM + h2) * IDIM + w2) - ((long)(d * IDIM + h) * IDIM + w);
                const float* p0 = x + idx;
                #pragma unroll
                for (int ic = 0; ic < CI; ++ic) {
                    const float xv = p0[(long)ic * IDIM * IDIM * IDIM];
                    #pragma unroll
                    for (int jc = 0; jc < CI; ++jc)
                        acc[ic * CI + jc] += xv * p0[(long)jc * IDIM * IDIM * IDIM + sh2];
                }
            }
        }
    }

    __shared__ float red[9][128];
    #pragma unroll
    for (int p = 0; p < CI * CI; ++p) red[p][threadIdx.x] = acc[p];
    __syncthreads();
    if (threadIdx.x < 9) {
        const int p = threadIdx.x;
        float s = 0.f;
        for (int i = 0; i < 128; ++i) s += red[p][i];
        atomicAdd(&C3[(long)off * 9 + p], (double)s);
    }
}

// ---------------------------------------------------------------------------
// miscCorr: C2[face][off2][ic][jc], C1[edge][off1][ic][jc], C0[corner][ic][jc]
// gridDim.y: 0..5 faces, 6..17 edges, 18 corners.
// ---------------------------------------------------------------------------
__global__ void miscCorr_kernel(const float* __restrict__ x, int nimg,
                                double* __restrict__ C2,
                                double* __restrict__ C1,
                                double* __restrict__ C0)
{
    const int g = blockIdx.y;
    float acc[CI * CI];
    #pragma unroll
    for (int i = 0; i < CI * CI; ++i) acc[i] = 0.f;

    if (g < 6) {
        const int fixed_ax = g / 2, side = g % 2;
        const int plane = side * 31;
        const int o1 = blockIdx.z / 3 - 1, o2 = blockIdx.z % 3 - 1;
        for (int n = 0; n < nimg; ++n) {
            const float* xb = x + (long)n * CI * IDIM * IDIM * IDIM;
            for (int p = threadIdx.x; p < 1024; p += blockDim.x) {
                const int u = p / 32, v = p % 32;
                const int u2 = u + o1, v2 = v + o2;
                if ((unsigned)u2 < 32u && (unsigned)v2 < 32u) {
                    int d, h, w, d2, h2, w2;
                    if (fixed_ax == 0)      { d = plane; h = u; w = v;  d2 = plane; h2 = u2; w2 = v2; }
                    else if (fixed_ax == 1) { h = plane; d = u; w = v;  h2 = plane; d2 = u2; w2 = v2; }
                    else                    { w = plane; d = u; h = v;  w2 = plane; d2 = u2; h2 = v2; }
                    const long o0 = ((long)(d * IDIM + h) * IDIM + w);
                    const long so = ((long)(d2 * IDIM + h2) * IDIM + w2) - o0;
                    #pragma unroll
                    for (int ic = 0; ic < CI; ++ic) {
                        const float xv = xb[(long)ic * IDIM * IDIM * IDIM + o0];
                        #pragma unroll
                        for (int jc = 0; jc < CI; ++jc)
                            acc[ic * CI + jc] += xv * xb[(long)jc * IDIM * IDIM * IDIM + o0 + so];
                    }
                }
            }
        }
        __shared__ float red[9][128];
        #pragma unroll
        for (int p = 0; p < CI * CI; ++p) red[p][threadIdx.x] = acc[p];
        __syncthreads();
        if (threadIdx.x < 9) {
            const int p = threadIdx.x;
            float s = 0.f;
            for (int i = 0; i < 128; ++i) s += red[p][i];
            atomicAdd(&C2[((long)g * 9 + blockIdx.z) * 9 + p], (double)s);
        }
    } else if (g < 18) {
        const int ge = g - 6;
        const int free_ax = ge / 4, sh2 = ge % 4;
        const int f1 = (sh2 >> 1) * 31, f2 = (sh2 & 1) * 31;
        const int o = (int)blockIdx.z - 1;
        for (int n = 0; n < nimg; ++n) {
            const float* xb = x + (long)n * CI * IDIM * IDIM * IDIM;
            if (threadIdx.x < 32) {
                const int u = threadIdx.x;
                const int u2 = u + o;
                if ((unsigned)u2 < 32u) {
                    int d, h, w, d2, h2, w2;
                    if (free_ax == 0)      { d = u;  h = f1; w = f2;  d2 = u2; h2 = f1; w2 = f2; }
                    else if (free_ax == 1) { h = u;  d = f1; w = f2;  h2 = u2; d2 = f1; w2 = f2; }
                    else                   { w = u;  d = f1; h = f2;  w2 = u2; d2 = f1; h2 = f2; }
                    const long o0 = ((long)(d * IDIM + h) * IDIM + w);
                    const long so = ((long)(d2 * IDIM + h2) * IDIM + w2) - o0;
                    #pragma unroll
                    for (int ic = 0; ic < CI; ++ic) {
                        const float xv = xb[(long)ic * IDIM * IDIM * IDIM + o0];
                        #pragma unroll
                        for (int jc = 0; jc < CI; ++jc)
                            acc[ic * CI + jc] += xv * xb[(long)jc * IDIM * IDIM * IDIM + o0 + so];
                    }
                }
            }
        }
        if (threadIdx.x < 32) {
            #pragma unroll
            for (int p = 0; p < CI * CI; ++p) {
                #pragma unroll
                for (int st = 16; st > 0; st >>= 1)
                    acc[p] += __shfl_down_sync(0xffffffff, acc[p], st);
                if (threadIdx.x == 0)
                    atomicAdd(&C1[((long)ge * 3 + blockIdx.z) * 9 + p], (double)acc[p]);
            }
        }
    } else {
        if (threadIdx.x < 8) {
            const int cn = threadIdx.x;
            const int d = ((cn >> 2) & 1) * 31, h = ((cn >> 1) & 1) * 31, w = (cn & 1) * 31;
            const long o0 = ((long)(d * IDIM + h) * IDIM + w);
            for (int n = 0; n < nimg; ++n) {
                const float* xb = x + (long)n * CI * IDIM * IDIM * IDIM;
                #pragma unroll
                for (int ic = 0; ic < CI; ++ic) {
                    const float xv = xb[(long)ic * IDIM * IDIM * IDIM + o0];
                    #pragma unroll
                    for (int jc = 0; jc < CI; ++jc)
                        acc[ic * CI + jc] += xv * xb[(long)jc * IDIM * IDIM * IDIM + o0];
                }
            }
            #pragma unroll
            for (int p = 0; p < CI * CI; ++p)
                atomicAdd(&C0[(long)cn * 9 + p], (double)acc[p]);
        }
    }
}

// ---------------------------------------------------------------------------
// bnFinalize: batch mean/var from XS + correlations, running-stat update,
// fused scale/shift (includes the 1/64 pool factor).
// ---------------------------------------------------------------------------
__global__ void bnFinalize_kernel(
    const double* __restrict__ XS,   // [81]
    const double* __restrict__ C3,   // [243]
    const double* __restrict__ C2,   // [486]
    const double* __restrict__ C1,   // [324]
    const double* __restrict__ C0,   // [72]
    const float* __restrict__ WT,    // [27][CI][CO]
    const float* __restrict__ M3,    // [27][CI][CI][CO]
    const float* __restrict__ M2,    // [6][9][CI][CI][CO]
    const float* __restrict__ M1,    // [12][3][CI][CI][CO]
    const float* __restrict__ M0,    // [8][CI][CI][CO]
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ cbias,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    long long* __restrict__ nbt,
    float* __restrict__ scale,
    float* __restrict__ shift,
    double m, double nimg, double momentum, double eps, int track)
{
    const int oc = blockIdx.x;   // 16 blocks, 64 threads
    double S3 = 0.0;
    for (int i = threadIdx.x; i < 243; i += 64) {
        const int off = i / 9, p = i % 9;
        S3 += (double)M3[(off * 9 + p) * CO + oc] * C3[i];
    }
    double S2 = 0.0;
    for (int i = threadIdx.x; i < 486; i += 64) {
       