```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: min over depth (dim=2) followed by softmax over channels (dim=1).
# The 3D convolution is left to cuDNN (highly optimized implicit GEMM), while the
# two reduction ops are fused into a single memory-bound kernel that reads the conv
# output exactly once and writes the final result directly.
fused_min_softmax_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cfloat>
#include <cstdint>

#define MAX_CHANNELS 32

// input  : (B, C, D, H, W) contiguous
// output : (B, C, H, W)    contiguous
// Each thread handles one (b, h, w) spatial location: it computes the min over
// depth for every channel (coalesced reads across threads), then performs the
// numerically-stable softmax over channels entirely in registers/local memory.
__global__ void min_depth_softmax_channel_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    const int C,
    const int D,
    const int HW,
    const int N)          // B * H * W
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N) return;

    const int b  = idx / HW;
    const int hw = idx - b * HW;

    const float* in_base = input + (int64_t)b * C * D * HW + hw;

    float mins[MAX_CHANNELS];

    // Min over the depth dimension for each channel.
    for (int c = 0; c < C; ++c) {
        const float* p = in_base + (int64_t)c * D * HW;
        float m = FLT_MAX;
        for (int d = 0; d < D; ++d) {
            m = fminf(m, __ldg(p + (int64_t)d * HW));
        }
        mins[c] = m;
    }

    // Softmax over the channel dimension (with max subtraction for stability).
    float mx = -FLT_MAX;
    for (int c = 0; c < C; ++c) {
        mx = fmaxf(mx, mins[c]);
    }

    float sum = 0.0f;
    for (int c = 0; c < C; ++c) {
        const float e = expf(mins[c] - mx);
        mins[c] = e;
        sum += e;
    }
    const float inv_sum = 1.0f / sum;

    float* out_base = output + (int64_t)b * C * HW + hw;
    for (int c = 0; c < C; ++c) {
        out_base[(int64_t)c * HW] = mins[c] * inv_sum;
    }
}

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

    auto xc = x.contiguous();
    const int B = xc.size(0);
    const int C = xc.size(1);
    const int D = xc.size(2);
    const int H = xc.size(3);
    const int W = xc.size(4);
    TORCH_CHECK(C <= MAX_CHANNELS, "channel count exceeds kernel limit");

    auto out = at::empty({B, C, H, W}, xc.options());

    const int HW = H * W;
    const int N = B * HW;
    if (N == 0) return out;

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

    auto stream = at::cuda::getCurrentCUDAStream();
    min_depth_softmax_channel_kernel<<<blocks, threads, 0, stream>>>(
        xc.data_ptr<float>(),
        out.data_ptr<float>(),
        C, D, HW, N);

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

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

fused_min_softmax = load_inline(
    name="fused_min_softmax",
    cpp_sources=fused_min_softmax_cpp_source,
    cuda_sources=fused_min_softmax_cuda_source,
    functions=["fused_min_softmax_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: keeps the cuDNN 3D convolution, but fuses the
    min-over-depth reduction and channel softmax into a single custom
    CUDA kernel (one pass over the data, no intermediate tensors).
    """
    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)
        """
        x = self.conv(x)
        # Fast path: min along depth (dim=2) + softmax over channels, fully fused.
        if (self.dim == 2 and x.is_cuda and x.dtype == torch.float32
                and x.size(1) <= 32):
            return fused_min_softmax.fused_min_softmax_cuda(x)
        # Generic fallback for any other configuration.
        x = torch.min(x, dim=self.dim)[0]
        x = torch.softmax(x, dim=1)
        return x
```