You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups. 

    You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.


        Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: 

        ``` 

        # 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


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, a, b):
        return a + b


def get_inputs():
    # randomly generate input tensors based on the model architecture
    a = torch.randn(1, 128).cuda()
    b = torch.randn(1, 128).cuda()
    return [a, b]


def get_init_inputs():
    # randomly generate tensors required for initialization based on the model architecture
    return []

        ``` 

        The example new arch with custom CUDA kernels looks like this:
        ```
        # 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

# Define the custom CUDA kernel for element-wise addition
elementwise_add_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void elementwise_add_kernel(const float* a, const float* b, float* out, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        out[idx] = a[idx] + b[idx];
    }
}

torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b) {
    auto size = a.numel();
    auto out = torch::zeros_like(a);

    const int block_size = 256;
    const int num_blocks = (size + block_size - 1) / block_size;

    elementwise_add_kernel<<<num_blocks, block_size>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size);

    return out;
}
"""

elementwise_add_cpp_source = (
    "torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b);"
)

# Compile the inline CUDA code for element-wise addition
elementwise_add = load_inline(
    name="elementwise_add",
    cpp_sources=elementwise_add_cpp_source,
    cuda_sources=elementwise_add_source,
    functions=["elementwise_add_cuda"],
    verbose=True,
    extra_cflags=[""],
    extra_ldflags=[""],
)


class ModelNew(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.elementwise_add = elementwise_add

    def forward(self, a, b):
        return self.elementwise_add.elementwise_add_cuda(a, b)

        ``` 

        
    You are given the following architecture: 

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

class Model(nn.Module):
    """
    Simple model that performs a 3D convolution, applies minimum operation along a specific dimension, 
    and then applies softmax.
    """
    def __init__(self, in_channels, out_channels, kernel_size, dim):
        super(Model, 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)
        x = torch.min(x, dim=self.dim)[0]  # Apply minimum along the specified dimension
        x = torch.softmax(x, dim=1)  # Apply softmax along the channel dimension
        return x

batch_size = 128
in_channels = 3
out_channels = 24  # Increased output channels
D, H, W = 24, 32, 32  # Increased depth
kernel_size = 3
dim = 2  # Dimension along which to apply minimum operation (e.g., depth)

def get_inputs():
    return [torch.rand(batch_size, in_channels, D, H, W)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, dim]
    ```
    
Optimize the architecture named Model with custom CUDA operators! Name your optimized output architecture ModelNew. Output the new code in codeblocks. Please generate real code, NOT pseudocode, make sure the code compiles and is fully functional. Just output the new model code, no other text, and NO testing code! 



This is turn 4 of 6 for this task. Your previous solution (turn 3) was:
```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

# ---------------------------------------------------------------------------
# Custom CUDA kernels
#
# 1) conv3d_min_softmax_tf32_kernel:
#    Fully fused kernel computing in ONE pass, with zero intermediates:
#       Conv3d (IC=3, K=3, COUT=24, stride 1, no padding)
#       -> min over depth (dim=2)
#       -> softmax over channels (dim=1)
#    The convolution is evaluated as an implicit GEMM on TF32 tensor cores
#    (mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32), which is allowed by
#    the task's precision rule (the PyTorch reference itself runs cuDNN with
#    TF32 enabled).  Each block handles 64 consecutive output positions
#    (flattened h*w) of one batch image; each of its 4 warps owns a 16x24
#    (position x channel) tile and sweeps all output depths, keeping a running
#    minimum in registers.  The 88-entry (81 real + 7 padding) patch-offset
#    table and the k-major weight matrix live in shared memory; weight
#    fragments are preloaded into registers once per block (they do not depend
#    on depth).  The channel softmax is a per-row reduction over the 4 lanes
#    of each accumulator quad (two __shfl_xor_sync steps), so no cross-warp
#    communication is ever needed.
#
# 2) min_softmax_kernel:
#    Fused min-over-depth + softmax-over-channels used on the cuDNN fallback
#    path (kept from the previous, known-correct solution).
# ---------------------------------------------------------------------------

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

#define DEVINL __device__ __forceinline__

DEVINL unsigned f2tf32(float x) {
    unsigned r;
    asm("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(r) : "f"(x));
    return r;
}

DEVINL void mma_tf32(float &d0, float &d1, float &d2, float &d3,
                     unsigned a0, unsigned a1, unsigned a2, unsigned a3,
                     unsigned b0, unsigned b1) {
    asm volatile(
        "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
        : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}

// x:  (B, 3, D, H, W), wt: (24, 3, 3, 3, 3), bias: (24)
// out: (B, 24, H-2, W-2)  == softmax_c min_d conv3d(x)
__global__ void __launch_bounds__(128)
conv3d_min_softmax_tf32_kernel(const float* __restrict__ x,
                               const float* __restrict__ wt,
                               const float* __restrict__ bias,
                               float* __restrict__ out,
                               int B, int D, int H, int W,
                               int H_OUT, int W_OUT) {
    constexpr int COUT = 24;
    constexpr int KP = 88;  // 81 real k + 7 zero padding

    __shared__ float sw[KP][COUT];
    __shared__ int   sk[KP];
    __shared__ float sb[COUT];

    const int tid = threadIdx.x;
    const int HW  = H * W;
    const int DHW = D * HW;

    for (int i = tid; i < KP * COUT; i += blockDim.x) {
        const int k = i / COUT, c = i % COUT;
        sw[k][c] = (k < 81) ? wt[c * 81 + k] : 0.f;
    }
    for (int i = tid; i < KP; i += blockDim.x) {
        int off = 0;
        if (i < 81) {
            const int ic = i / 27, rem = i % 27;
            const int kd = rem / 9, kh = (rem % 9) / 3, kw = rem % 3;
            off = ic * DHW + kd * HW + kh * W + kw;
        }
        sk[i] = off;
    }
    if (tid < COUT) sb[tid] = bias[tid];
    __syncthreads();

    const int HWo    = H_OUT * W_OUT;
    const int bpb    = (HWo + 63) >> 6;
    const int b      = blockIdx.x / bpb;
    const int p_base = (blockIdx.x - b * bpb) << 6;

    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int g    = lane >> 2;   // groupID
    const int t    = lane & 3;    // threadID_in_group
    const int wp0  = p_base + (warp << 4);

    const int rp0 = wp0 + g;
    const int rp1 = wp0 + g + 8;
    const bool v0 = rp0 < HWo;
    const bool v1 = rp1 < HWo;
    const int p0  = min(rp0, HWo - 1);
    const int p1  = min(rp1, HWo - 1);
    const int ro0 = (p0 / W_OUT) * W + (p0 % W_OUT);
    const int ro1 = (p1 / W_OUT) * W + (p1 % W_OUT);

    const float* xb = x + (long)b * 3 * DHW;

    // Preload weight fragments (independent of depth) into registers.
    unsigned breg[11][3][2];
    #pragma unroll
    for (int ks = 0; ks < 11; ++ks) {
        #pragma unroll
        for (int j = 0; j < 3; ++j) {
            breg[ks][j][0] = f2tf32(sw[ks * 8 + t][j * 8 + g]);
            breg[ks][j][1] = f2tf32(sw[ks * 8 + t + 4][j * 8 + g]);
        }
    }

    float minv[3][4];
    #pragma unroll
    for (int j = 0; j < 3; ++j)
        #pragma unroll
        for (int i = 0; i < 4; ++i) minv[j][i] = __int_as_float(0x7f800000);

    const int D_OUT = D - 2;
    for (int d = 0; d < D_OUT; ++d) {
        const float* xd = xb + (long)d * HW;
        float acc[3][4];
        #pragma unroll
        for (int j = 0; j < 3; ++j)
            #pragma unroll
            for (int i = 0; i < 4; ++i) acc[j][i] = 0.f;

        #pragma unroll
        for (int ks = 0; ks < 11; ++ks) {
            const int k0 = ks * 8 + t;
            const int k1 = k0 + 4;
            const unsigned a0 = f2tf32(__ldg(xd + ro0 + sk[k0]));
            const unsigned a1 = f2tf32(__ldg(xd + ro1 + sk[k0]));
            const unsigned a2 = f2tf32(__ldg(xd + ro0 + sk[k1]));
            const unsigned a3 = f2tf32(__ldg(xd + ro1 + sk[k1]));
            #pragma unroll
            for (int j = 0; j < 3; ++j) {
                mma_tf32(acc[j][0], acc[j][1], acc[j][2], acc[j][3],
                         a0, a1, a2, a3, breg[ks][j][0], breg[ks][j][1]);
            }
        }
        #pragma unroll
        for (int j = 0; j < 3; ++j)
            #pragma unroll
            for (int i = 0; i < 4; ++i) minv[j][i] = fminf(minv[j][i], acc[j][i]);
    }

    // bias (constant over depth, so it can be added after the min)
    float v[3][4];
    #pragma unroll
    for (int j = 0; j < 3; ++j) {
        const float bj0 = sb[j * 8 + 2 * t];
        const float bj1 = sb[j * 8 + 2 * t + 1];
        v[j][0] = minv[j][0] + bj0;
        v[j][1] = minv[j][1] + bj1;
        v[j][2] = minv[j][2] + bj0;
        v[j][3] = minv[j][3] + bj1;
    }

    // softmax over the 24 channels; each position's values live in a quad of
    // lanes (t = 0..3) -> reduce with two xor-shuffles.
    float m0 = fmaxf(fmaxf(v[0][0], v[0][1]), fmaxf(v[1][0], v[1][1]));
    m0 = fmaxf(m0, fmaxf(v[2][0], v[2][1]));
    m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, 1));
    m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, 2));
    float m1 = fmaxf(fmaxf(v[0][2], v[0][3]), fmaxf(v[1][2], v[1][3]));
    m1 = fmaxf(m1, fmaxf(v[2][2], v[2][3]));
    m1 = fmaxf(m1, __shfl_xor_sync(0xffffffffu, m1, 1));
    m1 = fmaxf(m1, __shfl_xor_sync(0xffffffffu, m1, 2));

    float s0 = 0.f, s1 = 0.f;
    #pragma unroll
    for (int j = 0; j < 3; ++j) {
        v[j][0] = expf(v[j][0] - m0); s0 += v[j][0];
        v[j][1] = expf(v[j][1] - m0); s0 += v[j][1];
        v[j][2] = expf(v[j][2] - m1); s1 += v[j][2];
        v[j][3] = expf(v[j][3] - m1); s1 += v[j][3];
    }
    s0 += __shfl_xor_sync(0xffffffffu, s0, 1);
    s0 += __shfl_xor_sync(0xffffffffu, s0, 2);
    s1 += __shfl_xor_sync(0xffffffffu, s1, 1);
    s1 += __shfl_xor_sync(0xffffffffu, s1, 2);
    const float r0 = 1.f / s0;
    const float r1 = 1.f / s1;

    const long ob = (long)b * COUT * HWo;
    #pragma unroll
    for (int j = 0; j < 3; ++j) {
        const int c0 = j * 8 + 2 * t;
        const int c1 = c0 + 1;
        if (v0) {
            out[ob + (long)c0 * HWo + rp0] = v[j][0] * r0;
            out[ob + (long)c1 * HWo + rp0] = v[j][1] * r0;
        }
        if (v1) {
            out[ob + (long)c0 * HWo + rp1] = v[j][2] * r1;
            out[ob + (long)c1 * HWo + rp1] = v[j][3] * r1;
        }
    }
}

// Fused min over depth -> softmax over channels (cuDNN fallback path).
// t: (B, 24, D2, H2, W2), out: (B, 24, H2, W2), W2 <= 32
__global__ void __launch_bounds__(256)
min_softmax_kernel(const float* __restrict__ t,
                   float* __restrict__ out,
                   int B, int D2, int H2, int W2) {
    constexpr int COUT = 24;
    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int gw   = blockIdx.x * (blockDim.x >> 5) + warp;
    const int rows = B * H2;
    if (gw >= rows) return;

    const int h      = gw % H2;
    const int b      = gw / H2;
    const int w      = min(lane, W2 - 1);
    const bool valid = (lane < W2);

    const long HW2 = (long)H2 * W2;
    const float* base = t + (long)b * COUT * D2 * HW2 + (long)h * W2 + w;

    float vals[COUT];
    #pragma unroll
    for (int c = 0; c < COUT; ++c) vals[c] = __ldg(base + (long)c * D2 * HW2);

    for (int d = 1; d < D2; ++d) {
        #pragma unroll
        for (int c = 0; c < COUT; ++c) {
            const float vv = __ldg(base + ((long)c * D2 + d) * HW2);
            vals[c] = fminf(vals[c], vv);
        }
    }

    float mx = vals[0];
    #pragma unroll
    for (int c = 1; c < COUT; ++c) mx = fmaxf(mx, vals[c]);
    float sum = 0.f;
    #pragma unroll
    for (int c = 0; c < COUT; ++c) { vals[c] = expf(vals[c] - mx); sum += vals[c]; }

    if (valid) {
        const long ob = (long)b * COUT * HW2 + (long)h * W2 + w;
        #pragma unroll
        for (int c = 0; c < COUT; ++c) out[ob + (long)c * HW2] = vals[c] / sum;
    }
}

torch::Tensor conv3d_min_softmax_tf32(torch::Tensor x, torch::Tensor weight,
                                      torch::Tensor bias) {
    const int B = x.size(0);
    const int D = x.size(2), H = x.size(3), W = x.size(4);
    const int D_OUT = D - 2, H_OUT = H - 2, W_OUT = W - 2;
    TORCH_CHECK(D_OUT >= 1 && H_OUT >= 1 && W_OUT >= 1, "bad shape");

    auto out = torch::empty({B, 24, H_OUT, W_OUT}, x.options());
    const long HWo = (long)H_OUT * W_OUT;
    const long bpb = (HWo + 63) / 64;
    const long grid = (long)B * bpb;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    conv3d_min_softmax_tf32_kernel<<<(int)grid, 128, 0, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
        out.data_ptr<float>(), B, D, H, W, H_OUT, W_OUT);
    return out;
}

torch::Tensor min_softmax(torch::Tensor t) {
    const int B  = t.size(0);
    const int D2 = t.size(2), H2 = t.size(3), W2 = t.size(4);
    TORCH_CHECK(W2 <= 32, "unsupported width");

    auto out = torch::empty({B, 24, H2, W2}, t.options());
    const int rows   = B * H2;
    const int wpb    = 8;
    const int blocks = (rows + wpb - 1) / wpb;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    min_softmax_kernel<<<blocks, wpb * 32, 0, stream>>>(
        t.data_ptr<float>(), out.data_ptr<float>(), B, D2, H2, W2);
    return out;
}
"""

cpp_source = r"""
torch::Tensor conv3d_min_softmax_tf32(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
torch::Tensor min_softmax(torch::Tensor t);
"""

_fused_ops = load_inline(
    name="fused_conv3d_min_softmax_tf32",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv3d_min_softmax_tf32", "min_softmax"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    Fast path: a single fused TF32 tensor-core CUDA kernel that performs
    Conv3d + min-over-depth + softmax-over-channels in one pass, writing only
    the final 4-D output (the large 5-D conv result is never materialized).
    Fallbacks: cuDNN Conv3d + fused min/softmax kernel, then pure PyTorch.
    The fastest correct path is selected once at runtime.
    """

    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._path = None  # 'tf32' or 'cudnn', chosen on first CUDA forward
        self._fast_ok = (
            in_channels == 3
            and out_channels == 24
            and kernel_size == 3
            and dim == 2
        )
        self._tf32_ok = None

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

    def _cudnn_path(self, x, w, b):
        return _fused_ops.min_softmax(F.conv3d(x, w, b))

    def _tf32_path(self, x, w, b):
        return _fused_ops.conv3d_min_softmax_tf32(x, w, b)

    def _supports_tf32(self):
        if self._tf32_ok is None:
            try:
                cap = torch.cuda.get_device_capability()
                self._tf32_ok = bool(cap[0] >= 8)
            except Exception:
                self._tf32_ok = False
        return self._tf32_ok

    def forward(self, x):
        if (
            not self._fast_ok
            or not x.is_cuda
            or x.dtype != torch.float32
            or x.dim() != 5
            or x.size(1) != 3
            or x.size(2) < 3
            or x.size(3) < 3
            or x.size(4) < 3
            or x.size(4) - 2 > 32  # cudnn fallback kernel needs W_OUT <= 32
        ):
            return self._reference(x)

        xc = x.contiguous()
        wc = self.conv.weight.contiguous()
        bc = self.conv.bias.contiguous()

        if self._path == "tf32":
            return self._tf32_path(xc, wc, bc)
        if self._path == "cudnn":
            return self._cudnn_path(xc, wc, bc)

        # One-time selection of the fastest correct path.
        r_cudnn = self._cudnn_path(xc, wc, bc)
        r_tf32 = None
        if self._supports_tf32():
            r_tf32 = self._tf32_path(xc, wc, bc)
            torch.cuda.synchronize()
            if not torch.allclose(r_tf32, r_cudnn, rtol=1e-2, atol=1e-2):
                r_tf32 = None  # fused kernel not trustworthy on this device

        if r_tf32 is None:
            self._path = "cudnn"
            return r_cudnn

        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        iters = 5

        start.record()
        for _ in range(iters):
            self._tf32_path(xc, wc, bc)
        end.record()
        torch.cuda.synchronize()
        t_tf32 = start.elapsed_time(end)

        start.record()
        for _ in range(iters):
            self._cudnn_path(xc, wc, bc)
        end.record()
        torch.cuda.synchronize()
        t_cudnn = start.elapsed_time(end)

        if t_tf32 <= t_cudnn:
            self._path = "tf32"
            return r_tf32
        self._path = "cudnn"
        return r_cudnn
```
Result of evaluating it:
The solution is correct on the test inputs. Mean runtime 0.322 ms; the PyTorch reference (with TF32 enabled) takes 2.06 ms; speedup 6.398x.
Profiler, top device kernels of one forward pass of your solution:
| device kernel | total device time (us) | calls |
|---|---|---|
| conv3d_min_softmax_tf32_kernel(float const*, float const*, float const*, float*, int, int, | 299.4 | 1 |

Write an improved complete ModelNew: fix any error, and if the solution is already correct, make it faster. Output the new code in one ```python codeblock. Just the model code, no other text, and NO testing code.


Precision rule: the timed forward must compute in FP32 (TF32 tensor-core math is allowed, as in the PyTorch reference). Do not cast inputs, weights or intermediate results to float16 or bfloat16 and do not call half-precision kernels in forward().
