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

# Fused CUDA kernel:
#   1) min-reduction over the depth dimension (dim 2) of the conv3d output
#   2) softmax over the channel dimension (dim 1) of the reduced tensor
# The conv3d output is viewed as (B, C, D, H*W). Each thread handles one
# spatial location hw of one batch element: it streams over (c, d) with
# fully coalesced reads, keeps the C minima in registers, then performs
# the channel softmax in registers and writes coalesced results.
# This removes an entire read+write pass over the large (B, C, D, H, W)
# tensor and a separate softmax kernel launch.

_cpp_source = "torch::Tensor fused_min_softmax_cuda(torch::Tensor x);"

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <float.h>

template <int C>
__global__ void fused_min_softmax_kernel(const float* __restrict__ x,
                                         float* __restrict__ out,
                                         int D, int HW) {
    const int hw = blockIdx.x * blockDim.x + threadIdx.x;
    if (hw >= HW) return;
    const int b = blockIdx.y;

    const float* xb = x + (size_t)b * C * D * HW;

    // Min over depth for every channel (minima kept in registers).
    float m[C];
    #pragma unroll
    for (int c = 0; c < C; ++c) {
        const float* p = xb + (size_t)c * D * HW + hw;
        float mn = p[0];
        #pragma unroll 4
        for (int d = 1; d < D; ++d) {
            mn = fminf(mn, p[(size_t)d * HW]);
        }
        m[c] = mn;
    }

    // Softmax over the channel dimension (numerically stabilized).
    float mx = m[0];
    #pragma unroll
    for (int c = 1; c < C; ++c) mx = fmaxf(mx, m[c]);

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

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

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

    const int64_t B = x.size(0);
    const int64_t C = x.size(1);
    const int64_t D = x.size(2);
    const int64_t H = x.size(3);
    const int64_t W = x.size(4);
    const int HW = static_cast<int>(H * W);

    auto out = at::empty({B, C, H, W}, x.options());
    if (out.numel() == 0) return out;

    const int block = 128;
    dim3 grid((HW + block - 1) / block, static_cast<unsigned int>(B));
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const float* xp = x.data_ptr<float>();
    float* op = out.data_ptr<float>();

#define LAUNCH_CASE(OC)                                                       \
    case OC:                                                                  \
        fused_min_softmax_kernel<OC><<<grid, block, 0, stream>>>(             \
            xp, op, static_cast<int>(D), HW);                                 \
        break;

    switch (C) {
        LAUNCH_CASE(1)  LAUNCH_CASE(2)  LAUNCH_CASE(3)  LAUNCH_CASE(4)
        LAUNCH_CASE(5)  LAUNCH_CASE(6)  LAUNCH_CASE(7)  LAUNCH_CASE(8)
        LAUNCH_CASE(9)  LAUNCH_CASE(10) LAUNCH_CASE(11) LAUNCH_CASE(12)
        LAUNCH_CASE(13) LAUNCH_CASE(14) LAUNCH_CASE(15) LAUNCH_CASE(16)
        LAUNCH_CASE(17) LAUNCH_CASE(18) LAUNCH_CASE(19) LAUNCH_CASE(20)
        LAUNCH_CASE(21) LAUNCH_CASE(22) LAUNCH_CASE(23) LAUNCH_CASE(24)
        LAUNCH_CASE(25) LAUNCH_CASE(26) LAUNCH_CASE(27) LAUNCH_CASE(28)
        LAUNCH_CASE(29) LAUNCH_CASE(30) LAUNCH_CASE(31) LAUNCH_CASE(32)
        LAUNCH_CASE(64)
        default:
            TORCH_CHECK(false, "unsupported channel count: ", C);
    }
#undef LAUNCH_CASE

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_ext = load_inline(
    name="fused_min_softmax_ext",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["fused_min_softmax_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)

_SUPPORTED_CHANNELS = set(range(1, 33)) | {64}


class ModelNew(nn.Module):
    """
    Optimized model: cuDNN 3D convolution followed by a single fused CUDA
    kernel that performs the min-reduction over the depth dimension and the
    softmax over the channel dimension in one pass.
    """
    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
        # Let cuDNN autotune the fastest conv algorithm for the fixed shapes.
        torch.backends.cudnn.benchmark = True

    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)
        """
        x = self.conv(x)
        dim = self.dim if self.dim >= 0 else self.dim + x.dim()
        if (dim == 2 and x.is_cuda and x.dtype == torch.float32
                and x.size(1) in _SUPPORTED_CHANNELS):
            return _ext.fused_min_softmax_cuda(x.contiguous())
        x = torch.min(x, dim=dim)[0]  # Apply minimum along the specified dimension
        x = torch.softmax(x, dim=1)   # Apply softmax along the channel dimension
        return x
```