```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 (3x3x3, stride 1, no padding) + min over depth
# (dim=2) + softmax over channels (dim=1), all in a single kernel launch.
#
# Strategy:
#  * One thread block per (batch, output-h) row of 30 output positions.
#  * The needed input tile (CIN x DIN x 3 x 32) and all conv weights are staged
#    into shared memory once per block.
#  * Each thread owns 3 output channels (24 channels / 8 thread-rows) at one
#    spatial position w, register-accumulates 11 depth slices at a time (weights
#    stay in registers across those 11 slices), and keeps a running per-channel
#    minimum across all 22 depth outputs.
#  * After a block-wide sync, the per-position minima (24 channels) are
#    normalized with a numerically stable softmax and written straight to the
#    output tensor. The intermediate conv output (95 MB) is never materialized.

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

#define CIN 3
#define COUT 24
#define KS 3
#define DIN 24
#define HIN 32
#define WIN 32
#define DOUT (DIN - KS + 1)   // 22
#define HOUT (HIN - KS + 1)   // 30
#define WOUT (WIN - KS + 1)   // 30
#define CGROUP 3              // output channels per thread
#define YDIM (COUT / CGROUP)  // 8
#define DBLK 11               // depth slices processed per register block
#define SMIN_STRIDE (COUT + 1)

__global__ __launch_bounds__(32 * YDIM)
void fused_conv3d_min_softmax_kernel(
    const float* __restrict__ x,      // (B, CIN, DIN, HIN, WIN)
    const float* __restrict__ wt,     // (COUT, CIN, KS, KS, KS)
    const float* __restrict__ bias,   // (COUT)
    float* __restrict__ out)          // (B, COUT, HOUT, WOUT)
{
    __shared__ float s_x[CIN * DIN * KS * WIN];
    __shared__ float s_w[COUT * CIN * KS * KS * KS];
    __shared__ float s_min[WOUT][SMIN_STRIDE];

    const int b  = blockIdx.x;
    const int h  = blockIdx.y;
    const int w  = threadIdx.x;   // 0..31 (30,31 inactive for compute)
    const int ty = threadIdx.y;   // 0..7  channel group
    const int tid = ty * 32 + w;
    const int nthreads = 32 * YDIM;

    const float* xb = x + (size_t)b * (CIN * DIN * HIN * WIN);

    // Stage the input tile x[b, :, :, h:h+3, :] into shared memory.
    #pragma unroll 4
    for (int i = tid; i < CIN * DIN * KS * WIN; i += nthreads) {
        const int chunk = i / (KS * WIN);   // ic*DIN + d
        const int rem   = i % (KS * WIN);   // kh*WIN + ww
        s_x[i] = xb[(size_t)chunk * (HIN * WIN) + (size_t)h * WIN + rem];
    }
    // Stage all conv weights.
    #pragma unroll 4
    for (int i = tid; i < COUT * CIN * KS * KS * KS; i += nthreads) {
        s_w[i] = wt[i];
    }
    __syncthreads();

    if (w < WOUT) {
        const int ch0 = ty * CGROUP;
        float biasv[CGROUP];
        #pragma unroll
        for (int c = 0; c < CGROUP; ++c) biasv[c] = bias[ch0 + c];

        float minv[CGROUP];
        #pragma unroll
        for (int c = 0; c < CGROUP; ++c) minv[c] = FLT_MAX;

        #pragma unroll
        for (int d0 = 0; d0 < DOUT; d0 += DBLK) {
            float acc[DBLK][CGROUP];
            #pragma unroll
            for (int db = 0; db < DBLK; ++db) {
                #pragma unroll
                for (int c = 0; c < CGROUP; ++c) acc[db][c] = biasv[c];
            }

            #pragma unroll
            for (int ic = 0; ic < CIN; ++ic) {
                #pragma unroll
                for (int kd = 0; kd < KS; ++kd) {
                    // Load the weights for this (ic, kd) slice into registers;
                    // they are reused across all DBLK depth slices.
                    float wr[CGROUP][KS * KS];
                    #pragma unroll
                    for (int c = 0; c < CGROUP; ++c) {
                        const float* wp = s_w + (((ch0 + c) * CIN + ic) * KS + kd) * (KS * KS);
                        #pragma unroll
                        for (int i = 0; i < KS * KS; ++i) wr[c][i] = wp[i];
                    }
                    const int xbase = ((ic * DIN + d0 + kd) * KS) * WIN + w;
                    #pragma unroll
                    for (int kh = 0; kh < KS; ++kh) {
                        #pragma unroll
                        for (int kw = 0; kw < KS; ++kw) {
                            #pragma unroll
                            for (int db = 0; db < DBLK; ++db) {
                                const float v = s_x[xbase + db * (KS * WIN) + kh * WIN + kw];
                                acc[db][0] = fmaf(v, wr[0][kh * KS + kw], acc[db][0]);
                                acc[db][1] = fmaf(v, wr[1][kh * KS + kw], acc[db][1]);
                                acc[db][2] = fmaf(v, wr[2][kh * KS + kw], acc[db][2]);
                            }
                        }
                    }
                }
            }
            #pragma unroll
            for (int db = 0; db < DBLK; ++db) {
                #pragma unroll
                for (int c = 0; c < CGROUP; ++c) minv[c] = fminf(minv[c], acc[db][c]);
            }
        }

        #pragma unroll
        for (int c = 0; c < CGROUP; ++c) s_min[w][ch0 + c] = minv[c];
    }
    __syncthreads();

    // Softmax over the channel dimension for each (b, h, w).
    if (ty == 0 && w < WOUT) {
        float m = -FLT_MAX;
        #pragma unroll
        for (int c = 0; c < COUT; ++c) m = fmaxf(m, s_min[w][c]);
        float ev[COUT];
        float sum = 0.f;
        #pragma unroll
        for (int c = 0; c < COUT; ++c) { ev[c] = expf(s_min[w][c] - m); sum += ev[c]; }
        const float inv = 1.f / sum;
        const size_t base = (((size_t)b * COUT) * HOUT + h) * WOUT + w;
        #pragma unroll
        for (int c = 0; c < COUT; ++c) out[base + (size_t)c * (HOUT * WOUT)] = ev[c] * inv;
    }
}

torch::Tensor fused_conv3d_min_softmax_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    const int64_t B = x.size(0);
    auto out = torch::empty({B, (int64_t)COUT, (int64_t)HOUT, (int64_t)WOUT}, x.options());
    dim3 grid((unsigned)B, (unsigned)HOUT);
    dim3 block(32, YDIM);
    auto stream = at::cuda::getCurrentCUDAStream();
    fused_conv3d_min_softmax_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>());
    return out;
}
"""

fused_conv3d_min_softmax_cpp_src = (
    "torch::Tensor fused_conv3d_min_softmax_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);"
)

fused_conv3d_min_softmax = load_inline(
    name="fused_conv3d_min_softmax",
    cpp_sources=fused_conv3d_min_softmax_cpp_src,
    cuda_sources=fused_conv3d_min_softmax_cuda_src,
    functions=["fused_conv3d_min_softmax_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: replaces Conv3d + min(dim) + softmax with a single fused
    custom CUDA kernel when the configuration matches the supported shape,
    and falls back to the original PyTorch ops otherwise.
    """
    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

        if isinstance(kernel_size, (tuple, list)):
            k = tuple(kernel_size)
        else:
            k = (kernel_size,) * 3
        ddim = dim if dim >= 0 else dim + 5  # normalize for 5D (N, C, D, H, W)
        self._use_fused = (
            in_channels == 3 and out_channels == 24 and k == (3, 3, 3)
            and ddim == 2
            and self.conv.bias is not None
            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
        )

    def forward(self, x):
        if (self._use_fused and 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):
            return fused_conv3d_min_softmax.fused_conv3d_min_softmax_cuda(
                x.contiguous(),
                self.conv.weight.contiguous(),
                self.conv.bias.contiguous(),
            )
        out = self.conv(x)
        out = torch.min(out, dim=self.dim)[0]
        out = torch.softmax(out, dim=1)
        return out
```