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

# Cache of compiled kernels, keyed by (in_channels, out_channels, kernel_size, D, H, W)
_module_cache = {}

_CUDA_TEMPLATE = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>

#define IC @IC@
#define OC @OC@
#define KSIZE @K@
#define ID @D@
#define IH @H@
#define IW @W@
#define OD (ID - KSIZE + 1)
#define OH (IH - KSIZE + 1)
#define OW (IW - KSIZE + 1)
#define POS_PER_BLOCK 8
#define NTHREADS (OC * POS_PER_BLOCK)
#define WSIZE (OC * IC * KSIZE * KSIZE * KSIZE)

// Fully fused kernel:
//   1) 3D convolution (stride=1, pad=0, dilation=1, groups=1)
//   2) minimum along the output depth dimension (dim=2)
//   3) softmax along the channel dimension (dim=1)
//
// Each thread owns one (batch, channel, h, w) output location, accumulates the
// convolution result for every output depth in registers (acc[OD]) while
// sliding over the input depth (each loaded value is reused for all K taps),
// tracks the running minimum over depth, then a block-level softmax over the
// OC channels is performed through shared memory.
__global__ __launch_bounds__(NTHREADS)
void fused_conv3d_min_softmax_kernel(
    const float* __restrict__ x,
    const float* __restrict__ wt,
    const float* __restrict__ bias,
    float* __restrict__ out,
    const int total_pos)
{
    __shared__ float ws[WSIZE];
    __shared__ float bs[OC];
    __shared__ float mins[OC][POS_PER_BLOCK];

    const int tid = threadIdx.x;
    for (int i = tid; i < WSIZE; i += NTHREADS) ws[i] = wt[i];
    if (tid < OC) bs[tid] = bias[tid];
    __syncthreads();

    // Consecutive thread ids map to consecutive spatial positions (coalesced
    // input reads and output writes); threads with the same position cover all
    // OC channels and read identical input windows (broadcast / L1 reuse).
    const int c = tid / POS_PER_BLOCK;   // output channel
    const int p = tid - c * POS_PER_BLOCK; // position inside the block
    const int pos = blockIdx.x * POS_PER_BLOCK + p;

    const int w0 = pos % OW;
    const int tmp = pos / OW;
    const int h0 = tmp % OH;
    const int b  = tmp / OH;
    const bool valid = (pos < total_pos);

    if (valid) {
        const float* xb = x + (long)b * (IC * ID * IH * IW);
        const float* wc = ws + c * (IC * KSIZE * KSIZE * KSIZE);

        float acc[OD];
        #pragma unroll
        for (int d = 0; d < OD; ++d) acc[d] = bs[c];

        #pragma unroll
        for (int ic = 0; ic < IC; ++ic) {
            const float* xic = xb + ic * (ID * IH * IW);
            #pragma unroll
            for (int kh = 0; kh < KSIZE; ++kh) {
                const int row = (h0 + kh) * IW + w0;
                #pragma unroll
                for (int kw = 0; kw < KSIZE; ++kw) {
                    const float* xz = xic + row + kw;
                    float wk[KSIZE];
                    #pragma unroll
                    for (int kd = 0; kd < KSIZE; ++kd)
                        wk[kd] = wc[((ic * KSIZE + kd) * KSIZE + kh) * KSIZE + kw];
                    // Slide along the input depth: each loaded value feeds all
                    // KSIZE depth taps of the convolution.
                    #pragma unroll
                    for (int t = 0; t < ID; ++t) {
                        const float z = __ldg(xz + t * (IH * IW));
                        #pragma unroll
                        for (int kd = 0; kd < KSIZE; ++kd) {
                            if (t >= kd && (t - kd) < OD) {
                                acc[t - kd] = fmaf(z, wk[kd], acc[t - kd]);
                            }
                        }
                    }
                }
            }
        }

        // Minimum over the depth dimension (dim=2).
        float m = acc[0];
        #pragma unroll
        for (int d = 1; d < OD; ++d) m = fminf(m, acc[d]);
        mins[c][p] = m;
    }
    __syncthreads();

    // Softmax over the channel dimension (dim=1) for this spatial position.
    if (valid) {
        float mx = mins[0][p];
        #pragma unroll
        for (int i = 1; i < OC; ++i) mx = fmaxf(mx, mins[i][p]);
        float sum = 0.0f;
        #pragma unroll
        for (int i = 0; i < OC; ++i) sum += __expf(mins[i][p] - mx);
        const float inv = 1.0f / sum;
        out[((long)b * OC + c) * (OH * OW) + h0 * OW + w0] =
            __expf(mins[c][p] - mx) * inv;
    }
}

torch::Tensor fused_conv3d_min_softmax(torch::Tensor x, torch::Tensor wt, torch::Tensor bias) {
    const int B = x.size(0);
    auto out = torch::empty({B, (long)OC, (long)OH, (long)OW}, x.options());
    const int total_pos = B * OH * OW;
    if (total_pos == 0) return out;
    const int blocks = (total_pos + POS_PER_BLOCK - 1) / POS_PER_BLOCK;
    fused_conv3d_min_softmax_kernel<<<blocks, NTHREADS>>>(
        x.data_ptr<float>(), wt.data_ptr<float>(), bias.data_ptr<float>(),
        out.data_ptr<float>(), total_pos);
    return out;
}
'''


def _build_module(in_channels, out_channels, kernel_size, D, H, W):
    src = _CUDA_TEMPLATE
    for token, val in (
        ("@IC@", in_channels),
        ("@OC@", out_channels),
        ("@K@", kernel_size),
        ("@D@", D),
        ("@H@", H),
        ("@W@", W),
    ):
        src = src.replace(token, str(val))
    name = "fused_conv3d_min_softmax_{}_{}_{}_{}_{}_{}".format(
        in_channels, out_channels, kernel_size, D, H, W
    )
    return load_inline(
        name=name,
        cpp_sources="torch::Tensor fused_conv3d_min_softmax(torch::Tensor x, torch::Tensor wt, torch::Tensor bias);",
        cuda_sources=src,
        functions=["fused_conv3d_min_softmax"],
        verbose=False,
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3"],
    )


class ModelNew(nn.Module):
    """
    Optimized version of Model: a single fused CUDA kernel performs the 3D
    convolution, the minimum along the depth dimension, and the softmax along
    the channel dimension. The kernel is compiled lazily (and specialized for)
    the actual input spatial dimensions, with all loop bounds as compile-time
    constants so the depth accumulators live in registers.
    """

    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.in_channels = in_channels
        self.out_channels = out_channels
        if isinstance(kernel_size, (tuple, list)):
            kernel_size = tuple(kernel_size)
            self._k = kernel_size[0] if all(int(v) == int(kernel_size[0]) for v in kernel_size) else None
        else:
            self._k = int(kernel_size)

    def _fallback(self, x):
        y = self.conv(x)
        y = torch.min(y, dim=self.dim)[0]
        y = torch.softmax(y, dim=1)
        return y

    def forward(self, x):
        conv = self.conv
        k = self._k
        try:
            dim = self.dim if self.dim >= 0 else self.dim + x.dim()
            fast = (
                k is not None
                and x.is_cuda
                and x.dtype == torch.float32
                and dim == 2
                and conv.groups == 1
                and conv.bias is not None
                and all(int(s) == 1 for s in conv.stride)
                and all(int(p) == 0 for p in conv.padding)
                and all(int(d) == 1 for d in conv.dilation)
                and x.size(2) >= k
                and x.size(3) >= k
                and x.size(4) >= k
            )
        except Exception:
            fast = False
        if not fast:
            return self._fallback(x)

        D, H, W = int(x.size(2)), int(x.size(3)), int(x.size(4))
        key = (self.in_channels, self.out_channels, k, D, H, W)
        mod = _module_cache.get(key)
        if mod is None:
            mod = _build_module(*key)
            _module_cache[key] = mod
        return mod.fused_conv3d_min_softmax(
            x.contiguous(), conv.weight.contiguous(), conv.bias.contiguous()
        )
```