```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)
#    The convolution is evaluated as an implicit GEMM 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; each of its 4 warps owns a 16x24
#    (position x channel) tile and sweeps all output depths, keeping a running
#    minimum in registers.  The 88-entry (81 real + 7 padding) patch-offset
#    table and the k-major weight matrix live in shared memory; weight
#    fragments are preloaded into registers once per block (they do not depend
#    on depth).  The channel softmax is a per-row reduction over the 4 lanes
#    of each accumulator quad (two __shfl_xor_sync steps), so no cross-warp
#    communication is ever needed.
#
# 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__

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), 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) {
    constexpr int COUT = 24;
    constexpr int KP = 88;  // 81 real k + 7 zero padding

    __shared__ float sw[KP][COUT];
    __shared__ int   sk[KP];
    __shared__ float sb[COUT];

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

    for (int i = tid; i < KP * COUT; i += blockDim.x) {
        const int k = i / COUT, c = i % COUT;
        sw[k][c] = (k < 81) ? wt[c * 81 + k] : 0.f;
    }
    for (int i = tid; i < KP; i += blockDim.x) {
        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 = ic * DHW + kd * HW + kh * W + kw;
        }
        sk[i] = off;
    }
    if (tid < COUT) sb[tid] = bias[tid];
    __syncthreads();

    const int HWo    = H_OUT * W_OUT;
    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
    const int t    = lane & 3;    // threadID_in_group
    const int wp0  = p_base + (warp << 4);

    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 ro0 = (p0 / W_OUT) * W + (p0 % W_OUT);
    const int ro1 = (p1 / W_OUT) * W + (p1 % W_OUT);

    const float* xb = x + (long)b * 3 * DHW;

    // Preload weight fragments (independent of depth) into registers.
    unsigned breg[11][3][2];
    #pragma unroll
    for (int ks = 0; ks < 11; ++ks) {
        #pragma unroll
        for (int j = 0; j < 3; ++j) {
            breg[ks][j][0] = f2tf32(sw[ks * 8 + t][j * 8 + g]);
            breg[ks][j][1] = f2tf32(sw[ks * 8 + t + 4][j * 8 + g]);
        }
    }

    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 float* xd = xb + (long)d * HW;
        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 unsigned a0 = f2tf32(__ldg(xd + ro0 + sk[k0]));
            const unsigned a1 = f2tf32(__ldg(xd + ro1 + sk[k0]));
            const unsigned a2 = f2tf32(__ldg(xd + ro0 + sk[k1]));
            const unsigned a3 = f2tf32(__ldg(xd + ro1 + sk[k1]));
            #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 (constant over depth, so it can be added after the min)
    float v[3][4];
    #pragma unroll
    for (int j = 0; j < 3; ++j) {
        const float bj0 = sb[j * 8 + 2 * t];
        const float bj1 = sb[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;
    }

    // softmax over the 24 channels; each position's values live in a quad of
    // lanes (t = 0..3) -> reduce with two xor-shuffles.
    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) {
    constexpr int COUT = 24;
    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;
    }
}

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");

    auto out = torch::empty({B, 24, H_OUT, W_OUT}, x.options());
    const long HWo = (long)H_OUT * W_OUT;
    const long bpb = (HWo + 63) / 64;
    const long grid = (long)B * bpb;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    conv3d_min_softmax_tf32_kernel<<<(int)grid, 128, 0, 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);
    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, 24, 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",
    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).
    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(3) < 3
            or x.size(4) < 3
            or x.size(4) - 2 > 32  # cudnn fallback kernel needs W_OUT <= 32
        ):
            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
```