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

# Fused CUDA kernel that replaces:
#   1) torch.min(x, dim=2)  -> reduction over the depth dimension
#   2) torch.softmax(x, dim=1) -> softmax over the channel dimension
# on the [B, C, D, H, W] output of the Conv3d, producing [B, C, H, W]
# in a single kernel launch with a single pass over the data.
fused_min_softmax_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <float.h>

// One thread block handles one (batch, h) slice.
// blockDim = (32, C): threadIdx.x indexes the width (W <= 32),
// threadIdx.y indexes the channel (C <= 32).
__global__ void fused_min_softmax_kernel(
        const float* __restrict__ x,   // [B, C, D, H, W]
        float* __restrict__ y,         // [B, C, H, W]
        int C, int D, int H, int W) {
    const int b = blockIdx.x / H;
    const int h = blockIdx.x % H;
    const int c = threadIdx.y;
    const int w = threadIdx.x;
    const long HW = (long)H * (long)W;

    __shared__ float s_min[32][32];

    // Phase 1: min over depth for this (b, c, h, w). Coalesced reads across w.
    if (w < W) {
        const float* ptr = x + (((long)b * C + c) * (long)D) * HW + (long)h * W + w;
        float m = ptr[0];
        #pragma unroll 4
        for (int d = 1; d < D; ++d) {
            m = fminf(m, ptr[(long)d * HW]);
        }
        s_min[c][w] = m;
    }
    __syncthreads();

    // Phase 2: softmax over channels for each w (numerically stable).
    // One warp (threadIdx.y == 0) handles all 32 width positions; writes are
    // coalesced because lanes iterate over channels in lockstep.
    if (threadIdx.y == 0 && w < W) {
        float mx = -FLT_MAX;
        for (int cc = 0; cc < C; ++cc) {
            mx = fmaxf(mx, s_min[cc][w]);
        }
        float sum = 0.0f;
        for (int cc = 0; cc < C; ++cc) {
            float e = expf(s_min[cc][w] - mx);
            s_min[cc][w] = e;
            sum += e;
        }
        const float inv = 1.0f / sum;
        float* out = y + ((long)b * C) * HW + (long)h * W + w;
        for (int cc = 0; cc < C; ++cc) {
            out[(long)cc * HW] = s_min[cc][w] * inv;
        }
    }
}

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

    auto xc = x.contiguous();
    const int B = (int)xc.size(0);
    const int C = (int)xc.size(1);
    const int D = (int)xc.size(2);
    const int H = (int)xc.size(3);
    const int W = (int)xc.size(4);
    TORCH_CHECK(C <= 32 && W <= 32, "fused fast path requires C<=32 and W<=32");

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

    dim3 block(32, C);
    dim3 grid(B * H);
    auto stream = at::cuda::getCurrentCUDAStream();
    fused_min_softmax_kernel<<<grid, block, 0, stream>>>(
        xc.data_ptr<float>(), y.data_ptr<float>(), C, D, H, W);
    return y;
}
"""

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

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


class ModelNew(nn.Module):
    """
    Optimized model: keeps the cuDNN Conv3d and replaces the
    (min over depth) + (softmax over channels) pair with a single
    fused custom CUDA kernel.
    """
    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.fused = fused_min_softmax
        # Let cuDNN autotune the best conv algorithm for this fixed shape.
        torch.backends.cudnn.benchmark = True

    def forward(self, x):
        x = self.conv(x)
        dim = self.dim if self.dim >= 0 else self.dim + x.dim()
        if (x.is_cuda and x.dtype == torch.float32 and dim == 2
                and x.size(1) <= 32 and x.size(4) <= 32):
            # Fused: min over dim=2 (depth), then softmax over dim=1 (channels).
            return self.fused.fused_min_softmax_cuda(x)
        # Fallback to the reference PyTorch ops for any other configuration.
        x = torch.min(x, dim=self.dim)[0]
        x = torch.softmax(x, dim=1)
        return x
```