```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_kernel:
#    Fully fused kernel that computes, in one pass:
#      Conv3d (no padding, stride 1)  ->  min over depth (dim=2)  ->  softmax
#      over channels (dim=1).
#    Each warp handles one (batch, h) output row (one lane per w position).
#    Input rows are streamed from global memory (L1/L2 resident), weights are
#    staged in shared memory, and two output channels are register-blocked
#    per pass to reuse each input value.  Softmax over the 24 channels is
#    done per-lane in registers, so no cross-thread communication is needed.
#
# 2) min_softmax_kernel:
#    Fused reduction kernel used on the cuDNN path: min over depth followed
#    by softmax over channels.
# ---------------------------------------------------------------------------

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

constexpr int IC    = 3;
constexpr int K     = 3;
constexpr int COUT  = 24;
constexpr int DOUT  = 22;
constexpr int PATCH = IC * K * K * K;   // 81

// Fused Conv3d -> min over depth -> softmax over channels.
// x:      (B, IC, D, H, W)          with D == DOUT + K - 1
// weight: (COUT, IC, K, K, K)
// bias:   (COUT)
// out:    (B, COUT, H_OUT, W_OUT)   with W_OUT <= 32
__global__ void __launch_bounds__(256)
conv3d_min_softmax_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) {
    __shared__ float s_w[COUT * PATCH];
    __shared__ float s_b[COUT];

    const int tid = threadIdx.x;
    for (int i = tid; i < COUT * PATCH; i += blockDim.x) s_w[i] = wt[i];
    if (tid < COUT) s_b[tid] = bias[tid];
    __syncthreads();

    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int gw   = blockIdx.x * (blockDim.x >> 5) + warp;
    const int rows = B * H_OUT;
    if (gw >= rows) return;

    const int h      = gw % H_OUT;
    const int b      = gw / H_OUT;
    const int w      = min(lane, W_OUT - 1);   // clamp so OOB lanes read safely
    const bool valid = (lane < W_OUT);

    const long HW  = (long)H * W;
    const float* xb = x + (long)b * IC * D * HW;

    float vals[COUT];

    #pragma unroll
    for (int cp = 0; cp < COUT; cp += 2) {
        float acc0[DOUT], acc1[DOUT];
        #pragma unroll
        for (int d = 0; d < DOUT; ++d) { acc0[d] = 0.f; acc1[d] = 0.f; }

        const float* w0 = s_w + cp * PATCH;
        const float* w1 = w0 + PATCH;

        for (int ic = 0; ic < IC; ++ic) {
            const float* xic  = xb + (long)ic * D * HW;
            const float* w0ic = w0 + ic * (K * K * K);
            const float* w1ic = w1 + ic * (K * K * K);
            for (int kd = 0; kd < K; ++kd) {
                const float* xkd  = xic + (long)kd * HW + (long)h * W + w;
                const float* w0kd = w0ic + kd * (K * K);
                const float* w1kd = w1ic + kd * (K * K);
                for (int kh = 0; kh < K; ++kh) {
                    const float* xkh  = xkd + (long)kh * W;
                    const float* w0kh = w0kd + kh * K;
                    const float* w1kh = w1kd + kh * K;
                    #pragma unroll
                    for (int kw = 0; kw < K; ++kw) {
                        const float wv0 = w0kh[kw];
                        const float wv1 = w1kh[kw];
                        const float* xp = xkh + kw;
                        #pragma unroll
                        for (int d = 0; d < DOUT; ++d) {
                            const float xv = __ldg(xp + (long)d * HW);
                            acc0[d] = fmaf(xv, wv0, acc0[d]);
                            acc1[d] = fmaf(xv, wv1, acc1[d]);
                        }
                    }
                }
            }
        }

        // min over depth, then add bias (constant over depth)
        float m0 = acc0[0], m1 = acc1[0];
        #pragma unroll
        for (int d = 1; d < DOUT; ++d) { m0 = fminf(m0, acc0[d]); m1 = fminf(m1, acc1[d]); }
        vals[cp]     = m0 + s_b[cp];
        vals[cp + 1] = m1 + s_b[cp + 1];
    }

    // softmax over channels (per lane, fully in registers)
    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 oHW = (long)H_OUT * W_OUT;
        const long ob  = (long)b * COUT * oHW + (long)h * W_OUT + w;
        #pragma unroll
        for (int c = 0; c < COUT; ++c) out[ob + (long)c * oHW] = vals[c] / sum;
    }
}

// Fused min over depth -> softmax over channels (for the cuDNN path).
// t:   (B, COUT, D2, H2, W2),  out: (B, COUT, 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;

    // initialize with the d = 0 slice, then reduce the rest
    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 v = __ldg(base + ((long)c * D2 + d) * HW2);
            vals[c] = fminf(vals[c], v);
        }
    }

    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(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 - K + 1, H_OUT = H - K + 1, W_OUT = W - K + 1;
    TORCH_CHECK(D_OUT == DOUT, "unsupported depth");
    TORCH_CHECK(W_OUT <= 32, "unsupported width");

    auto out = torch::empty({B, COUT, H_OUT, W_OUT}, x.options());
    const int rows   = B * H_OUT;
    const int wpb    = 8;                       // warps per block (256 threads)
    const int blocks = (rows + wpb - 1) / wpb;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    conv3d_min_softmax_kernel<<<blocks, wpb * 32, 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, 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(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",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv3d_min_softmax", "min_softmax"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


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

    Fast path: a single fused CUDA kernel performing Conv3d + min over depth +
    softmax over channels (all in FP32), or a cuDNN Conv3d followed by a fused
    min+softmax kernel.  The faster of the two is selected once at runtime.
    A pure PyTorch fallback keeps the module fully general.
    """

    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  # 'fused' or 'cudnn', chosen on first CUDA forward
        self._fast_ok = (
            in_channels == 3
            and out_channels == 24
            and kernel_size == 3
            and dim == 2
        )

    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 _fused_path(self, x, w, b):
        return _fused_ops.conv3d_min_softmax(x, w, b)

    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) != 24          # -> depth of conv output must be 22
            or x.size(4) - 2 > 32       # -> conv output width must fit a warp
            or x.size(3) < 3
            or x.size(4) < 3
        ):
            return self._reference(x)

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

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

        # One-time selection of the fastest path.
        r_fused = self._fused_path(xc, wc, bc)
        r_cudnn = self._cudnn_path(xc, wc, bc)
        torch.cuda.synchronize()

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

        start.record()
        for _ in range(iters):
            self._fused_path(xc, wc, bc)
        end.record()
        torch.cuda.synchronize()
        t_fused = 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_fused <= t_cudnn:
            self._path = "fused"
            return r_fused
        self._path = "cudnn"
        return r_cudnn
```