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

# Enable cuDNN autotuning (exhaustive) so the 3D convolution picks the fastest
# available algorithm for the fixed shapes used by this model, and allow TF32
# tensor-core execution (same numerics mode as the reference configuration).
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.matmul.allow_tf32 = True
try:
    torch.backends.cudnn.benchmark_limit = 0  # exhaustive autotuning
except Exception:
    pass

# Fused CUDA kernels implementing:
#   1) min over the depth dimension (dim=2) of the conv3d output (N, C, D, H, W)
#   2) softmax over the channel dimension (dim=1) of the result (N, C, H, W)
#
# Two layouts are supported:
#   * NHWC (channels_last_3d): reads the cuDNN NHWC conv output directly, so
#     no NHWC->NCHW transpose kernels are needed at all.
#   * NCHW (contiguous): fallback used whenever the conv output is not NHWC.
fused_min_softmax_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cmath>
#include <cstdint>

namespace {

// ---------------------------------------------------------------------------
// NHWC (channels_last_3d) fused min-over-depth + softmax-over-channels.
// One thread per (n, h, w) spatial location. The per-channel values live in
// registers (C is a compile-time constant). Loads are 128-bit (float4) since
// the C channel values are contiguous in NHWC; stores are coalesced across
// the warp because consecutive threads handle consecutive spatial positions.
// ---------------------------------------------------------------------------
template <int C>
__global__ __launch_bounds__(256) void fused_min_softmax_nhwc_kernel(
        const float* __restrict__ in,
        float* __restrict__ out,
        int N, int D, int HW) {
    constexpr int C4 = C / 4;
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N * HW) return;

    const int n = idx / HW;
    const int s = idx - n * HW;

    const float4* base = reinterpret_cast<const float4*>(in)
                       + ((size_t)n * (size_t)D * (size_t)HW + (size_t)s) * C4;
    const size_t dstride = (size_t)HW * C4;

    float m[C];
    {
        const float4* p = base;
        #pragma unroll
        for (int i = 0; i < C4; ++i) {
            const float4 v = p[i];
            m[4 * i + 0] = v.x;
            m[4 * i + 1] = v.y;
            m[4 * i + 2] = v.z;
            m[4 * i + 3] = v.w;
        }
    }
    for (int d = 1; d < D; ++d) {
        const float4* p = base + (size_t)d * dstride;
        #pragma unroll
        for (int i = 0; i < C4; ++i) {
            const float4 v = p[i];
            m[4 * i + 0] = fminf(m[4 * i + 0], v.x);
            m[4 * i + 1] = fminf(m[4 * i + 1], v.y);
            m[4 * i + 2] = fminf(m[4 * i + 2], v.z);
            m[4 * i + 3] = fminf(m[4 * i + 3], v.w);
        }
    }

    // numerically stable softmax over the channel values held in registers
    float maxv = m[0];
    #pragma unroll
    for (int c = 1; c < C; ++c) maxv = fmaxf(maxv, m[c]);

    float sum = 0.0f;
    #pragma unroll
    for (int c = 0; c < C; ++c) {
        m[c] = expf(m[c] - maxv);
        sum += m[c];
    }
    const float inv = 1.0f / sum;

    float* op = out + (size_t)n * C * HW + s;
    #pragma unroll
    for (int c = 0; c < C; ++c) {
        op[(size_t)c * HW] = m[c] * inv;
    }
}

// Generic NHWC fallback for arbitrary channel counts: one block per (n, h, w),
// threads cooperate over channels using shared memory.
__global__ void fused_min_softmax_nhwc_generic_kernel(
        const float* __restrict__ in,
        float* __restrict__ out,
        int N, int C, int D, int HW) {
    extern __shared__ float smem[];  // [0..C): per-channel mins/exps, then reduction scratch
    const int n = blockIdx.x / HW;
    const int s = blockIdx.x - n * HW;
    const int tid = threadIdx.x;
    const int nt = blockDim.x;

    const float* base = in + ((size_t)n * (size_t)D * (size_t)HW + (size_t)s) * (size_t)C;
    const size_t dstride = (size_t)HW * (size_t)C;

    for (int c = tid; c < C; c += nt) {
        float m = base[c];
        for (int d = 1; d < D; ++d) m = fminf(m, base[(size_t)d * dstride + c]);
        smem[c] = m;
    }
    __syncthreads();

    float* red = smem + C;
    float lmax = -INFINITY;
    for (int c = tid; c < C; c += nt) lmax = fmaxf(lmax, smem[c]);
    red[tid] = lmax;
    __syncthreads();
    for (int st = nt >> 1; st > 0; st >>= 1) {
        if (tid < st) red[tid] = fmaxf(red[tid], red[tid + st]);
        __syncthreads();
    }
    const float maxv = red[0];
    __syncthreads();

    float lsum = 0.0f;
    for (int c = tid; c < C; c += nt) {
        const float e = expf(smem[c] - maxv);
        smem[c] = e;
        lsum += e;
    }
    red[tid] = lsum;
    __syncthreads();
    for (int st = nt >> 1; st > 0; st >>= 1) {
        if (tid < st) red[tid] += red[tid + st];
        __syncthreads();
    }
    const float inv = 1.0f / red[0];

    float* op = out + (size_t)n * (size_t)C * (size_t)HW + s;
    for (int c = tid; c < C; c += nt) {
        op[(size_t)c * HW] = smem[c] * inv;
    }
}

// ---------------------------------------------------------------------------
// NCHW (contiguous) fused min-over-depth + softmax-over-channels (fallback).
// One thread per (n, h, w) spatial location; per-channel values in registers.
// ---------------------------------------------------------------------------
template <int C>
__global__ void fused_min_softmax_nchw_kernel(
        const float* __restrict__ in,
        float* __restrict__ out,
        int N, int D, int HW) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N * HW) return;

    const int n = idx / HW;
    const int s = idx - n * HW;

    const float* in_n = in + (size_t)n * C * D * HW + s;

    float vals[C];
    #pragma unroll
    for (int c = 0; c < C; ++c) {
        const float* p = in_n + (size_t)c * D * HW;
        float m = p[0];
        #pragma unroll 4
        for (int d = 1; d < D; ++d) {
            m = fminf(m, p[(size_t)d * HW]);
        }
        vals[c] = m;
    }

    float maxv = vals[0];
    #pragma unroll
    for (int c = 1; c < C; ++c) maxv = fmaxf(maxv, vals[c]);

    float sum = 0.0f;
    #pragma unroll
    for (int c = 0; c < C; ++c) {
        vals[c] = expf(vals[c] - maxv);
        sum += vals[c];
    }
    const float inv = 1.0f / sum;

    float* out_n = out + (size_t)n * C * HW + s;
    #pragma unroll
    for (int c = 0; c < C; ++c) {
        out_n[(size_t)c * HW] = vals[c] * inv;
    }
}

__global__ void fused_min_softmax_nchw_generic_kernel(
        const float* __restrict__ in,
        float* __restrict__ out,
        int N, int C, int D, int HW) {
    extern __shared__ float smem[];
    const int n = blockIdx.x / HW;
    const int s = blockIdx.x - n * HW;
    const int tid = threadIdx.x;
    const int nt = blockDim.x;

    const float* in_n = in + (size_t)n * C * D * HW + s;

    for (int c = tid; c < C; c += nt) {
        const float* p = in_n + (size_t)c * D * HW;
        float m = p[0];
        for (int d = 1; d < D; ++d) m = fminf(m, p[(size_t)d * HW]);
        smem[c] = m;
    }
    __syncthreads();

    float* red = smem + C;
    float lmax = -INFINITY;
    for (int c = tid; c < C; c += nt) lmax = fmaxf(lmax, smem[c]);
    red[tid] = lmax;
    __syncthreads();
    for (int st = nt >> 1; st > 0; st >>= 1) {
        if (tid < st) red[tid] = fmaxf(red[tid], red[tid + st]);
        __syncthreads();
    }
    const float maxv = red[0];
    __syncthreads();

    float lsum = 0.0f;
    for (int c = tid; c < C; c += nt) {
        const float e = expf(smem[c] - maxv);
        smem[c] = e;
        lsum += e;
    }
    red[tid] = lsum;
    __syncthreads();
    for (int st = nt >> 1; st > 0; st >>= 1) {
        if (tid < st) red[tid] += red[tid + st];
        __syncthreads();
    }
    const float inv = 1.0f / red[0];

    float* out_n = out + (size_t)n * C * HW + s;
    for (int c = tid; c < C; c += nt) {
        out_n[(size_t)c * HW] = smem[c] * inv;
    }
}

}  // namespace

torch::Tensor fused_min_softmax_nhwc_cuda(torch::Tensor input) {
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(input.dim() == 5, "input must be 5D (N, C, D, H, W)");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "input must be float32");
    TORCH_CHECK(input.is_contiguous(at::MemoryFormat::ChannelsLast3d),
                "input must be channels_last_3d contiguous");

    const int N = (int)input.size(0);
    const int C = (int)input.size(1);
    const int D = (int)input.size(2);
    const int H = (int)input.size(3);
    const int W = (int)input.size(4);
    TORCH_CHECK(D >= 1, "depth must be >= 1");

    auto out = at::empty({N, C, H, W}, input.options());
    const int HW = H * W;
    const long long total = (long long)N * HW;
    if (total == 0 || C == 0) return out;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* in_ptr = input.data_ptr<float>();
    float* out_ptr = out.data_ptr<float>();

    const bool aligned16 = ((reinterpret_cast<uintptr_t>(in_ptr)) & 0xF) == 0;
    const int threads = 256;
    const int blocks = (int)((total + threads - 1) / threads);

    bool launched = false;
    if (aligned16 && (C % 4) == 0) {
        #define FUSED_NHWC_CASE(CC)                                             \
            case CC:                                                            \
                fused_min_softmax_nhwc_kernel<CC><<<blocks, threads, 0, stream>>>( \
                    in_ptr, out_ptr, N, D, HW);                                 \
                launched = true;                                                \
                break;
        switch (C) {
            FUSED_NHWC_CASE(4)
            FUSED_NHWC_CASE(8)
            FUSED_NHWC_CASE(12)
            FUSED_NHWC_CASE(16)
            FUSED_NHWC_CASE(20)
            FUSED_NHWC_CASE(24)
            FUSED_NHWC_CASE(28)
            FUSED_NHWC_CASE(32)
            FUSED_NHWC_CASE(48)
            FUSED_NHWC_CASE(64)
            FUSED_NHWC_CASE(96)
            FUSED_NHWC_CASE(128)
            default:
                break;
        }
        #undef FUSED_NHWC_CASE
    }
    if (!launched) {
        TORCH_CHECK(C <= 8192, "channel count too large for generic kernel");
        const int gthreads = 128;
        const size_t shmem = ((size_t)C + (size_t)gthreads) * sizeof(float);
        fused_min_softmax_nhwc_generic_kernel<<<(int)total, gthreads, shmem, stream>>>(
            in_ptr, out_ptr, N, C, D, HW);
    }

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}

torch::Tensor fused_min_softmax_cuda(torch::Tensor input) {
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(input.dim() == 5, "input must be 5D (N, C, D, H, W)");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "input must be float32");

    auto in = input.contiguous();
    const int N = (int)in.size(0);
    const int C = (int)in.size(1);
    const int D = (int)in.size(2);
    const int H = (int)in.size(3);
    const int W = (int)in.size(4);
    TORCH_CHECK(D >= 1, "depth must be >= 1");

    auto out = at::empty({N, C, H, W}, in.options());
    const int HW = H * W;
    const long long total = (long long)N * HW;
    if (total == 0 || C == 0) return out;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* in_ptr = in.data_ptr<float>();
    float* out_ptr = out.data_ptr<float>();

    const int threads = 256;
    const int blocks = (int)((total + threads - 1) / threads);

    bool launched = false;
    #define FUSED_NCHW_CASE(CC)                                                  \
        case CC:                                                                 \
            fused_min_softmax_nchw_kernel<CC><<<blocks, threads, 0, stream>>>(   \
                in_ptr, out_ptr, N, D, HW);                                      \
            launched = true;                                                     \
            break;
    switch (C) {
        FUSED_NCHW_CASE(1)
        FUSED_NCHW_CASE(2)
        FUSED_NCHW_CASE(3)
        FUSED_NCHW_CASE(4)
        FUSED_NCHW_CASE(6)
        FUSED_NCHW_CASE(8)
        FUSED_NCHW_CASE(12)
        FUSED_NCHW_CASE(16)
        FUSED_NCHW_CASE(24)
        FUSED_NCHW_CASE(32)
        FUSED_NCHW_CASE(48)
        FUSED_NCHW_CASE(64)
        FUSED_NCHW_CASE(96)
        FUSED_NCHW_CASE(128)
        default:
            break;
    }
    #undef FUSED_NCHW_CASE
    if (!launched) {
        const int gthreads = 128;
        const size_t shmem = ((size_t)C + (size_t)gthreads) * sizeof(float);
        fused_min_softmax_nchw_generic_kernel<<<(int)total, gthreads, shmem, stream>>>(
            in_ptr, out_ptr, N, C, D, HW);
    }

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

fused_min_softmax_cpp_src = (
    "torch::Tensor fused_min_softmax_nhwc_cuda(torch::Tensor input);\n"
    "torch::Tensor fused_min_softmax_cuda(torch::Tensor input);"
)

_fused_ops = load_inline(
    name="fused_min_softmax_ext_v2",
    cpp_sources=fused_min_softmax_cpp_src,
    cuda_sources=fused_min_softmax_cuda_src,
    functions=["fused_min_softmax_nhwc_cuda", "fused_min_softmax_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model:
      - The input is converted once to channels_last_3d so cuDNN runs its NHWC
        tensor-core kernels natively, eliminating the NCHW<->NHWC transpose
        kernels (input, output, and the extra contiguous copy) that dominated
        the baseline runtime.
      - min-over-depth (dim=2) and softmax-over-channels (dim=1) are fused
        into a single custom CUDA kernel that reads the NHWC conv output
        directly with vectorized float4 loads, keeping all per-channel values
        in registers and never materializing the intermediate min tensor.
    """
    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

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, D, H, W)
        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, H, W)
        """
        dim = self.dim if self.dim >= 0 else self.dim + x.dim()
        fast = (dim == 2 and x.is_cuda and x.dtype == torch.float32)
        if fast:
            # One small transpose of the (128, 3, 24, 32, 32) input; lets cuDNN
            # use NHWC kernels end-to-end and returns a channels_last_3d output.
            x = x.contiguous(memory_format=torch.channels_last_3d)
        x = self.conv(x)
        if fast and x.is_contiguous(memory_format=torch.channels_last_3d):
            return _fused_ops.fused_min_softmax_nhwc_cuda(x)
        if fast and x.size(1) <= 4096:
            return _fused_ops.fused_min_softmax_cuda(x)
        x = torch.min(x, dim=dim)[0]
        x = torch.softmax(x, dim=1)
        return x
```