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

# Enable cuDNN autotuning (exhaustive) so the 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

# Custom CUDA kernels implementing:
#   1) a specialized NCHW -> NHWC transpose that also zero-pads the channel
#      dimension from C_in (1..4) to 4. cuDNN's NHWC tensor-core conv kernels
#      require the channel count to be a multiple of 4 to use vectorized
#      (float4) loads; with C_in=3 it falls back to a slow gather ("indexed")
#      kernel. The transpose is fully coalesced on both sides and needs no
#      shared memory: each thread owns one spatial position and all channels.
#   2) fused min-over-depth (dim=2) + softmax-over-channels (dim=1) that reads
#      the NHWC conv output directly (vectorized, streaming loads) and keeps
#      all per-channel values in registers.
#   3) an NCHW fallback variant of (2).
_fused_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 {

// ---------------------------------------------------------------------------
// NCHW -> NHWC transpose with channel zero-padding to 4 channels.
// in:  (N, C_IN, DS) contiguous
// out: (N, 4, DS) channels_last, i.e. (N, DS, 4) contiguous
// Each thread handles ILP spatial positions for all channels: the channel
// loads are coalesced across the warp and the float4 stores are coalesced
// across the warp, so no shared-memory staging is needed.
// ---------------------------------------------------------------------------
template <int C_IN>
__global__ __launch_bounds__(256) void nchw_to_nhwc4_kernel(
        const float* __restrict__ in,
        float* __restrict__ out,
        int DS) {
    constexpr int THREADS = 256;
    constexpr int ILP = 4;
    const int n = blockIdx.x;
    const int s0 = blockIdx.y * (THREADS * ILP);
    const int t = threadIdx.x;

    const float* in_n = in + (size_t)n * (size_t)C_IN * (size_t)DS;
    float4* out_n = reinterpret_cast<float4*>(out + (size_t)n * 4 * (size_t)DS);

    if (s0 + THREADS * ILP <= DS) {
        float v[C_IN][ILP];
        #pragma unroll
        for (int k = 0; k < ILP; ++k) {
            const int s = s0 + t + k * THREADS;
            #pragma unroll
            for (int c = 0; c < C_IN; ++c) {
                v[c][k] = in_n[(size_t)c * (size_t)DS + s];
            }
        }
        #pragma unroll
        for (int k = 0; k < ILP; ++k) {
            const int s = s0 + t + k * THREADS;
            float4 o;
            o.x = v[0][k];
            o.y = (C_IN > 1) ? v[1][k] : 0.0f;
            o.z = (C_IN > 2) ? v[2][k] : 0.0f;
            o.w = (C_IN > 3) ? v[3][k] : 0.0f;
            out_n[s] = o;
        }
    } else {
        float v[C_IN][ILP];
        bool valid[ILP];
        #pragma unroll
        for (int k = 0; k < ILP; ++k) {
            const int s = s0 + t + k * THREADS;
            valid[k] = (s < DS);
            if (valid[k]) {
                #pragma unroll
                for (int c = 0; c < C_IN; ++c) {
                    v[c][k] = in_n[(size_t)c * (size_t)DS + s];
                }
            }
        }
        #pragma unroll
        for (int k = 0; k < ILP; ++k) {
            const int s = s0 + t + k * THREADS;
            if (valid[k]) {
                float4 o;
                o.x = v[0][k];
                o.y = (C_IN > 1) ? v[1][k] : 0.0f;
                o.z = (C_IN > 2) ? v[2][k] : 0.0f;
                o.w = (C_IN > 3) ? v[3][k] : 0.0f;
                out_n[s] = o;
            }
        }
    }
}

// ---------------------------------------------------------------------------
// 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) and
// streaming (ldcs); stores are coalesced across the warp.
// ---------------------------------------------------------------------------
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 = __ldcs(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;
        }
    }
    int d = 1;
    for (; d + 1 < D; d += 2) {
        const float4* p0 = base + (size_t)d * dstride;
        const float4* p1 = p0 + dstride;
        #pragma unroll
        for (int i = 0; i < C4; ++i) {
            const float4 v0 = __ldcs(p0 + i);
            const float4 v1 = __ldcs(p1 + i);
            m[4 * i + 0] = fminf(m[4 * i + 0], fminf(v0.x, v1.x));
            m[4 * i + 1] = fminf(m[4 * i + 1], fminf(v0.y, v1.y));
            m[4 * i + 2] = fminf(m[4 * i + 2], fminf(v0.z, v1.z));
            m[4 * i + 3] = fminf(m[4 * i + 3], fminf(v0.w, v1.w));
        }
    }
    if (d < D) {
        const float4* p = base + (size_t)d * dstride;
        #pragma unroll
        for (int i = 0; i < C4; ++i) {
            const float4 v = __ldcs(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).
// ---------------------------------------------------------------------------
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 nchw_to_nhwc_pad4_cuda(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(x.dim() == 5, "input must be 5D (N, C, D, H, W)");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "input must be float32");
    TORCH_CHECK(x.is_contiguous(), "input must be NCHW contiguous");

    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int D = (int)x.size(2);
    const int H = (int)x.size(3);
    const int W = (int)x.size(4);
    TORCH_CHECK(C >= 1 && C <= 4, "channel count must be in [1, 4]");

    auto out = at::empty({N, 4, D, H, W},
                         x.options().memory_format(at::MemoryFormat::ChannelsLast3d));
    const long long DS = (long long)D * H * W;
    if (DS == 0 || N == 0) return out;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    constexpr int THREADS = 256;
    constexpr int ILP = 4;
    const dim3 grid(N, (unsigned)((DS + THREADS * ILP - 1) / (THREADS * ILP)));
    const float* in_ptr = x.data_ptr<float>();
    float* out_ptr = out.data_ptr<float>();

    switch (C) {
        case 1:
            nchw_to_nhwc4_kernel<1><<<grid, THREADS, 0, stream>>>(in_ptr, out_ptr, (int)DS);
            break;
        case 2:
            nchw_to_nhwc4_kernel<2><<<grid, THREADS, 0, stream>>>(in_ptr, out_ptr, (int)DS);
            break;
        case 3:
            nchw_to_nhwc4_kernel<3><<<grid, THREADS, 0, stream>>>(in_ptr, out_ptr, (int)DS);
            break;
        default:
            nchw_to_nhwc4_kernel<4><<<grid, THREADS, 0, stream>>>(in_ptr, out_ptr, (int)DS);
            break;
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}

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_cpp_src = (
    "torch::Tensor nchw_to_nhwc_pad4_cuda(torch::Tensor x);\n"
    "torch::Tensor fused_min_softmax_nhwc_cuda(torch::Tensor input);\n"
    "torch::Tensor fused_min_softmax_cuda(torch::Tensor input);"
)

_ext = load_inline(
    name="fused_conv_min_softmax_ext_v3",
    cpp_sources=_fused_cpp_src,
    cuda_sources=_fused_cuda_src,
    functions=[
        "nchw_to_nhwc_pad4_cuda",
        "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:
      - A custom CUDA kernel transposes the input to channels_last_3d AND
        zero-pads the channel dimension 3 -> 4 in a single fully-coalesced
        pass. cuDNN's NHWC tensor-core kernels need the channel count to be a
        multiple of 4 to use vectorized loads; with C=3 it picks a slow
        gather kernel, so the padded input unlocks the fast implicit-GEMM
        path (the padded channel is all zeros and the corresponding padded
        conv weights are zero, so the result is bit-identical math).
      - min-over-depth (dim=2) and softmax-over-channels (dim=1) remain fused
        into a single custom CUDA kernel that reads the NHWC conv output
        directly with streaming float4 loads, keeps all per-channel values in
        registers, and never materializes 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
        self._w4_cache = None
        self._w4_version = -1

    def _padded_weight(self):
        w = self.conv.weight
        wc = self._w4_cache
        if wc is None or self._w4_version != w._version or wc.device != w.device:
            # (O, I, kD, kH, kW) -> pad input-channel dim 3 -> 4 with zeros
            wc = F.pad(w, (0, 0, 0, 0, 0, 0, 0, 1)).contiguous()
            self._w4_cache = wc
            self._w4_version = w._version
        return wc

    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)
        y = None
        if fast and self.conv.groups == 1:
            if x.size(1) >= 1 and x.size(1) <= 4 and x.is_contiguous():
                # Fast path: fused transpose + zero-pad to 4 channels (NHWC),
                # cuDNN NHWC tensor-core conv, then fused min+softmax.
                x4 = _ext.nchw_to_nhwc_pad4_cuda(x)
                w4 = self.conv.weight if x.size(1) == 4 else self._padded_weight()
                y = F.conv3d(x4, w4, self.conv.bias, self.conv.stride,
                             self.conv.padding, self.conv.dilation,
                             self.conv.groups)
            else:
                y = self.conv(x.contiguous(memory_format=torch.channels_last_3d))
            if y.is_contiguous(memory_format=torch.channels_last_3d):
                return _ext.fused_min_softmax_nhwc_cuda(y)
            if y.size(1) <= 4096:
                return _ext.fused_min_softmax_cuda(y)
            x = y
        else:
            x = self.conv(x)
        x = torch.min(x, dim=dim)[0]
        x = torch.softmax(x, dim=1)
        return x
```