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

# Fully fused CUDA kernel:
#   Conv3d (3->24 ch, 3x3x3) + min over depth (dim=2) + softmax over channels (dim=1)
# One thread block handles one (batch, h) slice: it stages the required input
# tile and all weights in shared memory, computes every conv output along the
# depth dimension with register-blocked sliding windows, keeps only the running
# minimum over depth (the full 243MB conv tensor is never materialized), and
# finishes with a numerically stable softmax across channels.
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <float.h>

#define IN_C 3
#define OUT_C 24
#define KSIZE 3
#define IN_D 24
#define IN_H 32
#define IN_W 32
#define OUT_D (IN_D - KSIZE + 1)   // 22
#define OUT_H (IN_H - KSIZE + 1)   // 30
#define OUT_W (IN_W - KSIZE + 1)   // 30

__global__ void fused_conv3d_min_softmax_kernel(
    const float* __restrict__ x,
    const float* __restrict__ wt,
    const float* __restrict__ bs,
    float* __restrict__ out)
{
    __shared__ float xs[IN_C][IN_D][KSIZE][IN_W];
    __shared__ float ws[OUT_C * IN_C * KSIZE * KSIZE * KSIZE];
    __shared__ float bias_sh[OUT_C];
    __shared__ float mins[OUT_C][OUT_W];

    const int h   = blockIdx.x % OUT_H;
    const int b   = blockIdx.x / OUT_H;
    const int tid = threadIdx.x;

    const float* xb = x + (long long)b * (IN_C * IN_D * IN_H * IN_W);

    // Stage the input tile x[b, :, 0:IN_D, h:h+KSIZE, :] into shared memory.
    for (int idx = tid; idx < IN_C * IN_D * KSIZE * IN_W; idx += blockDim.x) {
        int w  = idx % IN_W;
        int t  = idx / IN_W;
        int kh = t % KSIZE;
        t /= KSIZE;
        int d  = t % IN_D;
        int ci = t / IN_D;
        xs[ci][d][kh][w] = xb[((ci * IN_D + d) * IN_H + (h + kh)) * IN_W + w];
    }
    // Stage all weights (layout matches the conv weight tensor exactly).
    for (int idx = tid; idx < OUT_C * IN_C * KSIZE * KSIZE * KSIZE; idx += blockDim.x)
        ws[idx] = wt[idx];
    if (tid < OUT_C) bias_sh[tid] = bs[tid];
    __syncthreads();

    // Each thread computes a few (channel, w) pairs: full conv for every depth,
    // keeping only the running minimum over depth.
    for (int p = tid; p < OUT_C * OUT_W; p += blockDim.x) {
        const int c = p / OUT_W;
        const int w = p % OUT_W;

        float acc[OUT_D];
#pragma unroll
        for (int d = 0; d < OUT_D; ++d) acc[d] = 0.f;

        const float* wp = ws + c * (IN_C * KSIZE * KSIZE * KSIZE);
        for (int ci = 0; ci < IN_C; ++ci) {
            for (int kh = 0; kh < KSIZE; ++kh) {
                for (int kw = 0; kw < KSIZE; ++kw) {
                    float xv[IN_D];
#pragma unroll
                    for (int dd = 0; dd < IN_D; ++dd)
                        xv[dd] = xs[ci][dd][kh][w + kw];
                    const float w0 = wp[((ci * KSIZE + 0) * KSIZE + kh) * KSIZE + kw];
                    const float w1 = wp[((ci * KSIZE + 1) * KSIZE + kh) * KSIZE + kw];
                    const float w2 = wp[((ci * KSIZE + 2) * KSIZE + kh) * KSIZE + kw];
#pragma unroll
                    for (int d = 0; d < OUT_D; ++d) {
                        acc[d] = fmaf(xv[d],     w0, acc[d]);
                        acc[d] = fmaf(xv[d + 1], w1, acc[d]);
                        acc[d] = fmaf(xv[d + 2], w2, acc[d]);
                    }
                }
            }
        }
        const float bias = bias_sh[c];
        float m = acc[0] + bias;
#pragma unroll
        for (int d = 1; d < OUT_D; ++d) m = fminf(m, acc[d] + bias);
        mins[c][w] = m;
    }
    __syncthreads();

    // Softmax over the channel dimension for each w of this (b, h) slice.
    if (tid < OUT_W) {
        const int w = tid;
        float mx = -FLT_MAX;
#pragma unroll
        for (int c = 0; c < OUT_C; ++c) mx = fmaxf(mx, mins[c][w]);
        float s = 0.f;
#pragma unroll
        for (int c = 0; c < OUT_C; ++c) {
            const float e = expf(mins[c][w] - mx);
            mins[c][w] = e;
            s += e;
        }
        const float inv = 1.f / s;
        float* outp = out + ((long long)b * OUT_C * OUT_H + h) * OUT_W + w;
#pragma unroll
        for (int c = 0; c < OUT_C; ++c)
            outp[c * (OUT_H * OUT_W)] = mins[c][w] * inv;
    }
}

torch::Tensor fused_conv3d_min_softmax(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.dim() == 5 && x.size(1) == IN_C && x.size(2) == IN_D &&
                x.size(3) == IN_H && x.size(4) == IN_W, "unexpected input shape");
    TORCH_CHECK(weight.numel() == OUT_C * IN_C * KSIZE * KSIZE * KSIZE, "unexpected weight shape");
    TORCH_CHECK(bias.numel() == OUT_C, "unexpected bias shape");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = bias.contiguous();

    const int B = (int)xc.size(0);
    auto out = at::empty({B, OUT_C, OUT_H, OUT_W}, xc.options());

    dim3 grid(B * OUT_H);
    dim3 block(256);
    fused_conv3d_min_softmax_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bc.data_ptr<float>(), out.data_ptr<float>());
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor fused_conv3d_min_softmax(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);"
)

fused_op = load_inline(
    name="fused_conv3d_min_softmax",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["fused_conv3d_min_softmax"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: a single fused CUDA kernel performs the 3D convolution,
    the minimum reduction along the depth dimension, and the softmax over the
    channel dimension, without ever materializing the full convolution output.
    Falls back to the native PyTorch ops for any configuration outside the
    specialized fast path.
    """
    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_op

    def _fast_path_ok(self, x):
        return (
            x.is_cuda
            and x.dtype == torch.float32
            and x.dim() == 5
            and x.size(1) == 3
            and x.size(2) == 24
            and x.size(3) == 32
            and x.size(4) == 32
            and self.conv.out_channels == 24
            and tuple(self.conv.kernel_size) == (3, 3, 3)
            and tuple(self.conv.stride) == (1, 1, 1)
            and tuple(self.conv.padding) == (0, 0, 0)
            and tuple(self.conv.dilation) == (1, 1, 1)
            and self.conv.groups == 1
            and self.conv.bias is not None
            and self.dim == 2
        )

    def forward(self, x):
        if self._fast_path_ok(x):
            return self.fused.fused_conv3d_min_softmax(x, self.conv.weight, self.conv.bias)
        out = self.conv(x)
        out = torch.min(out, dim=self.dim)[0]
        out = torch.softmax(out, dim=1)
        return out
```