```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
#
# 1) conv3d_min_softmax_tf32_kernel:
#    Fully fused kernel computing in ONE pass, with zero intermediates:
#       Conv3d (IC=3, K=3, COUT=24, stride 1, no padding)
#       -> min over depth (dim=2)
#       -> softmax over channels (dim=1)
#
#    Optimizations over the previous revision:
#    - The convolution runs on TF32 tensor cores
#      (mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32), which is allowed
#      by the task's precision rule (the PyTorch reference itself runs cuDNN
#      with TF32 enabled).
#    - Each block handles 64 consecutive output positions (flattened h*w) of
#      one batch image and sweeps every output depth.  Before the sweep it
#      stages the whole required input region (<= 6 input rows x 32 w x 3
#      channels x 24 depths, <= 56 KB) into dynamic shared memory with fully
#      coalesced 128-bit loads, converting values to TF32 *once* during the
#      staging.  The inner loop therefore only issues low-latency shared
#      loads (no global traffic, no cvt instructions), and the depth loop
#      needs no __syncthreads at all.
#    - Weight fragments are preloaded into registers once per block (they do
#      not depend on position or depth) via a small shared staging buffer
#      that is afterwards overwritten by the input staging (shared-memory
#      aliasing keeps the footprint at ~56 KB so 2 blocks fit per SM).
#    - Running depth-minimum stays in registers; the channel softmax is a
#      per-quad reduction (two __shfl_xor_sync), so no cross-warp traffic.
#
# 2) min_softmax_kernel:
#    Fused min-over-depth + softmax-over-channels used on the cuDNN fallback
#    path (kept from the previous, known-correct solution).
# ---------------------------------------------------------------------------

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

#define DEVINL __device__ __forceinline__

namespace {

constexpr int COUT   = 24;
constexpr int KP     = 88;    // 81 real reduction terms + 7 zero padding
constexpr int NR     = 6;     // staged input rows (worst case for 64 positions)
constexpr int IC     = 3;
constexpr int WMAX   = 32;
constexpr int DMAX   = 32;
constexpr int SD     = NR * IC * WMAX;                 // floats per depth slab
constexpr int STAGE  = DMAX * SD;                      // 18432 floats (72 KB)
constexpr int LDSW   = (COUT * 81 * 4 + 15) / 16;      // float4s of weights

DEVINL unsigned f2tf32(float x) {
    unsigned r;
    asm("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(r) : "f"(x));
    return r;
}

DEVINL void mma_tf32(float &d0, float &d1, float &d2, float &d3,
                     unsigned a0, unsigned a1, unsigned a2, unsigned a3,
                     unsigned b0, unsigned b1) {
    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"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}

// x:  (B, 3, D, H, W) with W <= 32, D <= 32
// wt: (24, 3, 3, 3, 3), bias: (24)
// out: (B, 24, H-2, W-2) == softmax_c min_d conv3d(x)
__global__ void __launch_bounds__(128)
conv3d_min_softmax_tf32_kernel(const float* __restrict__ x,
                               const float* __restrict__ wt,
                               const float* __restrict__ bias,
                               float* __restrict__ out,
                               int B, int D, int H, int W,
                               int H_OUT, int W_OUT, int HWo) {
    extern __shared__ float s_raw[];   // STAGE floats (weights staged first, then input)

    __shared__ int   sk[KP];
    __shared__ float sbias[COUT];

    const int tid  = threadIdx.x;
    const int HW   = H * W;
    const int DHW  = D * HW;

    const int bpb    = (HWo + 63) >> 6;
    const int b      = blockIdx.x / bpb;
    const int p_base = (blockIdx.x - b * bpb) << 6;

    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int g    = lane >> 2;   // groupID within mma tile
    const int t    = lane & 3;    // threadID-in-group
    const int wp0  = p_base + (warp << 4);

    // ---- per-thread position data -------------------------------------
    const int rp0 = wp0 + g;
    const int rp1 = wp0 + g + 8;
    const bool v0 = rp0 < HWo;
    const bool v1 = rp1 < HWo;
    const int p0  = min(rp0, HWo - 1);
    const int p1  = min(rp1, HWo - 1);

    const int h_first = p_base / W_OUT;
    const int h_last  = min(p_base + 63, HWo - 1) / W_OUT;
    const int hmax    = h_last - h_first;   // rows beyond this are never read

    // local row offsets inside the staged region
    const int ro0 = (p0 / W_OUT - h_first) * W + (p0 % W_OUT);
    const int ro1 = (p1 / W_OUT - h_first) * W + (p1 % W_OUT);

    // ---- stage weights -> shared (tf32), then into registers ----------
    const int4* w4 = reinterpret_cast<const int4*>(wt);
    int4* s4 = reinterpret_cast<int4*>(s_raw);
    for (int i = tid; i < LDSW; i += 128) {
        const int4 pk = w4[i];
        int4 o;
        o.x = (int)f2tf32(__int_as_float(pk.x));
        o.y = (int)f2tf32(__int_as_float(pk.y));
        o.z = (int)f2tf32(__int_as_float(pk.z));
        o.w = (int)f2tf32(__int_as_float(pk.w));
        s4[i] = o;
    }
    for (int i = tid; i < KP; i += 128) {
        int off = 0;
        if (i < 81) {
            const int ic = i / 27, rem = i % 27;
            const int kd = rem / 9, kh = (rem % 9) / 3, kw = rem % 3;
            off = kd * SD + ic * (NR * W) + kh * W + kw;
        }
        sk[i] = off;
    }
    if (tid < COUT) sbias[tid] = bias[tid];
    __syncthreads();

    unsigned breg[11][3][2];
    {
        const unsigned* sw = reinterpret_cast<const unsigned*>(s_raw);
        #pragma unroll
        for (int ks = 0; ks < 11; ++ks) {
            #pragma unroll
            for (int j = 0; j < 3; ++j) {
                breg[ks][j][0] = sw[(j * 8 + g) * 81 + ks * 8 + t];
                breg[ks][j][1] = sw[(j * 8 + g) * 81 + ks * 8 + t + 4];
            }
        }
    }
    __syncthreads();   // everyone done reading the weight staging area

    // ---- stage the input region (tf32) over the same shared memory ----
    {
        const float* xb = x + (long)b * IC * DHW;
        const int sd4   = SD >> 2;
        for (int r = 0; r < NR; ++r) {
            if (r > hmax) break;
            const float* srcb = xb + (long)min(h_first + r, H - 1) * W;
            float* dstb = s_raw + r * W;
            for (int e = tid; e < D * IC * sd4; e += 128) {
                const int d  = e / (IC * sd4);
                const int rm = e - d * (IC * sd4);
                const int ic = rm / sd4;
                const int wo = (rm - ic * sd4) << 2;
                const float4 pk = *reinterpret_cast<const float4*>(
                    srcb + (long)ic * DHW + (long)d * HW + wo);
                float4 o;
                o.x = __uint_as_float(f2tf32(pk.x));
                o.y = __uint_as_float(f2tf32(pk.y));
                o.z = __uint_as_float(f2tf32(pk.z));
                o.w = __uint_as_float(f2tf32(pk.w));
                *reinterpret_cast<float4*>(dstb + d * SD + ic * (NR * W) + wo) = o;
            }
        }
    }
    __syncthreads();

    // ---- main sweep over output depth: conv + running min -------------
    const unsigned* s_u = reinterpret_cast<const unsigned*>(s_raw);

    float minv[3][4];
    #pragma unroll
    for (int j = 0; j < 3; ++j)
        #pragma unroll
        for (int i = 0; i < 4; ++i) minv[j][i] = __int_as_float(0x7f800000);

    const int D_OUT = D - 2;
    for (int d = 0; d < D_OUT; ++d) {
        const unsigned* a0p = s_u + d * SD + ro0;
        const unsigned* a1p = s_u + d * SD + ro1;
        float acc[3][4];
        #pragma unroll
        for (int j = 0; j < 3; ++j)
            #pragma unroll
            for (int i = 0; i < 4; ++i) acc[j][i] = 0.f;

        #pragma unroll
        for (int ks = 0; ks < 11; ++ks) {
            const int k0 = ks * 8 + t;
            const int k1 = k0 + 4;
            const int o0 = sk[k0];
            const int o1 = sk[k1];
            const unsigned a0 = a0p[o0];
            const unsigned a1 = a1p[o0];
            const unsigned a2 = a0p[o1];
            const unsigned a3 = a1p[o1];
            #pragma unroll
            for (int j = 0; j < 3; ++j) {
                mma_tf32(acc[j][0], acc[j][1], acc[j][2], acc[j][3],
                         a0, a1, a2, a3, breg[ks][j][0], breg[ks][j][1]);
            }
        }
        #pragma unroll
        for (int j = 0; j < 3; ++j)
            #pragma unroll
            for (int i = 0; i < 4; ++i) minv[j][i] = fminf(minv[j][i], acc[j][i]);
    }

    // ---- bias + softmax over the 24 channels --------------------------
    float v[3][4];
    #pragma unroll
    for (int j = 0; j < 3; ++j) {
        const float bj0 = sbias[j * 8 + 2 * t];
        const float bj1 = sbias[j * 8 + 2 * t + 1];
        v[j][0] = minv[j][0] + bj0;
        v[j][1] = minv[j][1] + bj1;
        v[j][2] = minv[j][2] + bj0;
        v[j][3] = minv[j][3] + bj1;
    }

    // each position's 24 channel values live in a quad of lanes (t = 0..3)
    float m0 = fmaxf(fmaxf(v[0][0], v[0][1]), fmaxf(v[1][0], v[1][1]));
    m0 = fmaxf(m0, fmaxf(v[2][0], v[2][1]));
    m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, 1));
    m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, 2));
    float m1 = fmaxf(fmaxf(v[0][2], v[0][3]), fmaxf(v[1][2], v[1][3]));
    m1 = fmaxf(m1, fmaxf(v[2][2], v[2][3]));
    m1 = fmaxf(m1, __shfl_xor_sync(0xffffffffu, m1, 1));
    m1 = fmaxf(m1, __shfl_xor_sync(0xffffffffu, m1, 2));

    float s0 = 0.f, s1 = 0.f;
    #pragma unroll
    for (int j = 0; j < 3; ++j) {
        v[j][0] = expf(v[j][0] - m0); s0 += v[j][0];
        v[j][1] = expf(v[j][1] - m0); s0 += v[j][1];
        v[j][2] = expf(v[j][2] - m1); s1 += v[j][2];
        v[j][3] = expf(v[j][3] - m1); s1 += v[j][3];
    }
    s0 += __shfl_xor_sync(0xffffffffu, s0, 1);
    s0 += __shfl_xor_sync(0xffffffffu, s0, 2);
    s1 += __shfl_xor_sync(0xffffffffu, s1, 1);
    s1 += __shfl_xor_sync(0xffffffffu, s1, 2);
    const float r0 = 1.f / s0;
    const float r1 = 1.f / s1;

    const long ob = (long)b * COUT * HWo;
    #pragma unroll
    for (int j = 0; j < 3; ++j) {
        const int c0 = j * 8 + 2 * t;
        const int c1 = c0 + 1;
        if (v0) {
            out[ob + (long)c0 * HWo + rp0] = v[j][0] * r0;
            out[ob + (long)c1 * HWo + rp0] = v[j][1] * r0;
        }
        if (v1) {
            out[ob + (long)c0 * HWo + rp1] = v[j][2] * r1;
            out[ob + (long)c1 * HWo + rp1] = v[j][3] * r1;
        }
    }
}

// Fused min over depth -> softmax over channels (cuDNN fallback path).
// t: (B, 24, D2, H2, W2), out: (B, 24, H2, W2), W2 <= 32
__global__ void __launch_bounds__(256)
min_softmax_kernel(const float* __restrict__ t,
                   float* __restrict__ out,
                   int B, int D2, int H2, int W2) {
    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int gw   = blockIdx.x * (blockDim.x >> 5) + warp;
    const int rows = B * H2;
    if (gw >= rows) return;

    const int h      = gw % H2;
    const int b      = gw / H2;
    const int w      = min(lane, W2 - 1);
    const bool valid = (lane < W2);

    const long HW2 = (long)H2 * W2;
    const float* base = t + (long)b * COUT * D2 * HW2 + (long)h * W2 + w;

    float vals[COUT];
    #pragma unroll
    for (int c = 0; c < COUT; ++c) vals[c] = __ldg(base + (long)c * D2 * HW2);

    for (int d = 1; d < D2; ++d) {
        #pragma unroll
        for (int c = 0; c < COUT; ++c) {
            const float vv = __ldg(base + ((long)c * D2 + d) * HW2);
            vals[c] = fminf(vals[c], vv);
        }
    }

    float mx = vals[0];
    #pragma unroll
    for (int c = 1; c < COUT; ++c) mx = fmaxf(mx, vals[c]);
    float sum = 0.f;
    #pragma unroll
    for (int c = 0; c < COUT; ++c) { vals[c] = expf(vals[c] - mx); sum += vals[c]; }

    if (valid) {
        const long ob = (long)b * COUT * HW2 + (long)h * W2 + w;
        #pragma unroll
        for (int c = 0; c < COUT; ++c) out[ob + (long)c * HW2] = vals[c] / sum;
    }
}

}  // namespace

torch::Tensor conv3d_min_softmax_tf32(torch::Tensor x, torch::Tensor weight,
                                      torch::Tensor bias) {
    const int B = x.size(0);
    const int D = x.size(2), H = x.size(3), W = x.size(4);
    const int D_OUT = D - 2, H_OUT = H - 2, W_OUT = W - 2;
    TORCH_CHECK(D_OUT >= 1 && H_OUT >= 1 && W_OUT >= 1, "bad shape");
    TORCH_CHECK(W <= WMAX && D <= DMAX, "shape exceeds fused kernel limits");

    auto out = torch::empty({B, COUT, H_OUT, W_OUT}, x.options());
    const int HWo  = H_OUT * W_OUT;
    const long bpb = (HWo + 63) / 64;
    const long grid = (long)B * bpb;

    static bool configured = false;
    if (!configured) {
        cudaFuncSetAttribute(conv3d_min_softmax_tf32_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             STAGE * 4);
        cudaFuncSetAttribute(conv3d_min_softmax_tf32_kernel,
                             cudaFuncAttributePreferredSharedMemoryCarveout,
                             100);
        configured = true;
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    conv3d_min_softmax_tf32_kernel<<<(int)grid, 128, STAGE * 4, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
        out.data_ptr<float>(), B, D, H, W, H_OUT, W_OUT, HWo);
    return out;
}

torch::Tensor min_softmax(torch::Tensor t) {
    const int B  = t.size(0);
    const int D2 = t.size(2), H2 = t.size(3), W2 = t.size(4);
    TORCH_CHECK(W2 <= 32, "unsupported width");

    auto out = torch::empty({B, COUT, H2, W2}, t.options());
    const int rows   = B * H2;
    const int wpb    = 8;
    const int blocks = (rows + wpb - 1) / wpb;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    min_softmax_kernel<<<blocks, wpb * 32, 0, stream>>>(
        t.data_ptr<float>(), out.data_ptr<float>(), B, D2, H2, W2);
    return out;
}
"""

cpp_source = r"""
torch::Tensor conv3d_min_softmax_tf32(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
torch::Tensor min_softmax(torch::Tensor t);
"""

_fused_ops = load_inline(
    name="fused_conv3d_min_softmax_tf32_v2",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv3d_min_softmax_tf32", "min_softmax"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    Fast path: a single fused TF32 tensor-core CUDA kernel that performs
    Conv3d + min-over-depth + softmax-over-channels in one pass, writing only
    the final 4-D output (the large 5-D conv result is never materialized).
    The kernel stages the needed input region in shared memory once per block
    (pre-converted to TF32) so the depth sweep runs entirely out of shared
    memory with weight fragments held in registers.
    Fallbacks: cuDNN Conv3d + fused min/softmax kernel, then pure PyTorch.
    The fastest correct path is selected once at runtime.
    """

    def __init__(self, in_channels, out_channels, kernel_size, dim):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.dim = dim
        self._path = None  # 'tf32' or 'cudnn', chosen on first CUDA forward
        self._fast_ok = (
            in_channels == 3
            and out_channels == 24
            and kernel_size == 3
            and dim == 2
        )
        self._tf32_ok = None

    def _reference(self, x):
        x = self.conv(x)
        x = torch.min(x, dim=self.dim)[0]
        x = torch.softmax(x, dim=1)
        return x

    def _cudnn_path(self, x, w, b):
        return _fused_ops.min_softmax(F.conv3d(x, w, b))

    def _tf32_path(self, x, w, b):
        return _fused_ops.conv3d_min_softmax_tf32(x, w, b)

    def _supports_tf32(self):
        if self._tf32_ok is None:
            try:
                cap = torch.cuda.get_device_capability()
                self._tf32_ok = bool(cap[0] >= 8)
            except Exception:
                self._tf32_ok = False
        return self._tf32_ok

    def forward(self, x):
        if (
            not self._fast_ok
            or not x.is_cuda
            or x.dtype != torch.float32
            or x.dim() != 5
            or x.size(1) != 3
            or x.size(2) < 3
            or x.size(2) > 32
            or x.size(3) < 3
            or x.size(4) < 3
            or x.size(4) > 32  # fused kernels need W <= 32 (W_OUT <= 30)
        ):
            return self._reference(x)

        xc = x.contiguous()
        wc = self.conv.weight.contiguous()
        bc = self.conv.bias.contiguous()

        if self._path == "tf32":
            return self._tf32_path(xc, wc, bc)
        if self._path == "cudnn":
            return self._cudnn_path(xc, wc, bc)

        # One-time selection of the fastest correct path.
        r_cudnn = self._cudnn_path(xc, wc, bc)
        r_tf32 = None
        if self._supports_tf32():
            r_tf32 = self._tf32_path(xc, wc, bc)
            torch.cuda.synchronize()
            if not torch.allclose(r_tf32, r_cudnn, rtol=1e-2, atol=1e-2):
                r_tf32 = None  # fused kernel not trustworthy on this device

        if r_tf32 is None:
            self._path = "cudnn"
            return r_cudnn

        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        iters = 5

        start.record()
        for _ in range(iters):
            self._tf32_path(xc, wc, bc)
        end.record()
        torch.cuda.synchronize()
        t_tf32 = start.elapsed_time(end)

        start.record()
        for _ in range(iters):
            self._cudnn_path(xc, wc, bc)
        end.record()
        torch.cuda.synchronize()
        t_cudnn = start.elapsed_time(end)

        if t_tf32 <= t_cudnn:
            self._path = "tf32"
            return r_tf32
        self._path = "cudnn"
        return r_cudnn
```