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):
    """
    Performs a depthwise 2D convolution with asymmetric input and asymmetric kernel.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size_h (int): Height of the convolution kernel.
        kernel_size_w (int): Width of the convolution kernel.
        stride_h (int, optional): Stride of the convolution in height dimension. Defaults to 1.
        stride_w (int, optional): Stride of the convolution in width dimension. Defaults to 1.
        padding_h (int, optional): Padding applied to the input in height dimension. Defaults to 0.
        padding_w (int, optional): Padding applied to the input in width dimension. Defaults to 0.
        dilation_h (int, optional): Spacing between kernel elements in height dimension. Defaults to 1.
        dilation_w (int, optional): Spacing between kernel elements in width dimension. Defaults to 1.
        groups (int, optional): Number of blocked connections from input channels to output channels. Defaults to 1.
        bias (bool, optional): If `True`, adds a learnable bias to the output. Defaults to `False`.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size_h: int, kernel_size_w: int, stride_h: int = 1, stride_w: int = 1, padding_h: int = 0, padding_w: int = 0, dilation_h: int = 1, dilation_w: int = 1, groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, in_channels, (kernel_size_h, kernel_size_w), stride=(stride_h, stride_w), padding=(padding_h, padding_w), dilation=(dilation_h, dilation_w), groups=in_channels, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the depthwise 2D convolution.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, height_out, width_out).
        """
        return self.conv2d(x)

# Test code
batch_size = 32
in_channels = 128
out_channels = 128
kernel_size_h = 3
kernel_size_w = 7
width = 256
height = 128
stride_h = 1
stride_w = 1
padding_h = 0
padding_w = 0
dilation_h = 1
dilation_w = 1
groups = in_channels

def get_inputs():
    x = torch.rand(batch_size, in_channels, height, width)
    return [x]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size_h, kernel_size_w, stride_h, stride_w, padding_h, padding_w, dilation_h, dilation_w, groups]
    ```
    
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

_depthwise_conv2d_cpp_src = (
    "#include <torch/extension.h>\n"
    "torch::Tensor depthwise_conv2d_forward("
    "const torch::Tensor& input, "
    "const torch::Tensor& weight, "
    "const c10::optional<torch::Tensor>& bias, "
    "const int64_t stride_h, const int64_t stride_w, "
    "const int64_t pad_h, const int64_t pad_w, "
    "const int64_t dil_h, const int64_t dil_w);"
)

_depthwise_conv2d_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cstdint>

#define TILE_W 32
#define TILE_H 16

// ---------------------------------------------------------------------------
// Vectorized shared-memory load helpers (C++11-compatible template structs).
// LdsA: pointer must be 16B aligned. LdsB: pointer must be 8B aligned.
// ---------------------------------------------------------------------------
template <int N>
struct LdsA {
    static __device__ __forceinline__ void run(const float* __restrict__ p, float* __restrict__ v) {
        *reinterpret_cast<float4*>(v) = *reinterpret_cast<const float4*>(p);
        LdsA<N - 4>::run(p + 4, v + 4);
    }
};
template <>
struct LdsA<0> {
    static __device__ __forceinline__ void run(const float*, float*) {}
};
template <>
struct LdsA<2> {
    static __device__ __forceinline__ void run(const float* __restrict__ p, float* __restrict__ v) {
        *reinterpret_cast<float2*>(v) = *reinterpret_cast<const float2*>(p);
    }
};

template <int N>
struct LdsB {
    static __device__ __forceinline__ void run(const float* __restrict__ p, float* __restrict__ v) {
        *reinterpret_cast<float2*>(v) = *reinterpret_cast<const float2*>(p);
        LdsB<N - 2>::run(p + 2, v + 2);
    }
};
template <>
struct LdsB<0> {
    static __device__ __forceinline__ void run(const float*, float*) {}
};
template <>
struct LdsB<1> {
    static __device__ __forceinline__ void run(const float* __restrict__ p, float* __restrict__ v) {
        v[0] = p[0];
    }
};

// Window loader: pick widest safe vector width based on per-thread output
// block width OW (OW==4 => column offset is a multiple of 4 floats => 16B).
template <int OW, int W>
struct WinLoader {
    static __device__ __forceinline__ void run(const float* p, float* v) {
        LdsB<W>::run(p, v);
    }
};
template <int W>
struct WinLoader<4, W> {
    static __device__ __forceinline__ void run(const float* p, float* v) {
        LdsA<W>::run(p, v);
    }
};

// Vectorized global-store helper; p must be N*4-byte aligned.
template <int N>
struct StgV {
    static __device__ __forceinline__ void run(float* __restrict__ p, const float* __restrict__ v) {
        *reinterpret_cast<float4*>(p) = *reinterpret_cast<const float4*>(v);
        StgV<N - 4>::run(p + 4, v + 4);
    }
};
template <>
struct StgV<0> {
    static __device__ __forceinline__ void run(float*, const float*) {}
};
template <>
struct StgV<2> {
    static __device__ __forceinline__ void run(float* __restrict__ p, const float* __restrict__ v) {
        *reinterpret_cast<float2*>(p) = *reinterpret_cast<const float2*>(v);
    }
};

// ---------------------------------------------------------------------------
// Register-blocked kernel (float, stride == 1, dilation == 1).
// Each block computes a (32*OW) x (8*OH) output tile; each thread computes
// OW x OH outputs. Weights are cached in registers and the per-thread input
// window is loaded once from shared memory with wide vector loads (sliding
// window reuse across all of the thread's outputs), drastically cutting
// shared-memory traffic versus the one-output-per-thread kernel.
// Grid: (ceil(out_w/(32*OW)), ceil(out_h/(8*OH)), batch*channels)
// Block: (32, 8)
// ---------------------------------------------------------------------------
template <int KH, int KW, int OW, int OH>
__global__ __launch_bounds__(256, 3) void depthwise_conv2d_reg_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    const int channels,
    const int in_h, const int in_w,
    const int out_h, const int out_w,
    const int pad_h, const int pad_w,
    const bool has_bias)
{
    constexpr int TX = 32;
    constexpr int TY = 8;
    constexpr int TILE_OW = TX * OW;
    constexpr int TILE_OH = TY * OH;
    constexpr int TILE_IW = TILE_OW + KW - 1;
    constexpr int TILE_IH = TILE_OH + KH - 1;
    constexpr int TILE_IWP = (TILE_IW + 3) & ~3;   // pad rows to 16B multiples
    constexpr int W = KW + OW - 1;                 // window registers per row
    constexpr int VH = KH + OH - 1;                // window register rows
    constexpr int NT = TX * TY;
    constexpr int NLOAD = TILE_IH * TILE_IW;

    __shared__ float s_w[KH * KW];
    __shared__ __align__(16) float s_in[TILE_IH][TILE_IWP];

    const long bc = blockIdx.z;
    const int c = (int)(bc % channels);
    const int tid = threadIdx.y * TX + threadIdx.x;

    if (tid < KH * KW) {
        s_w[tid] = weight[(long)c * KH * KW + tid];
    }

    const int ow0 = blockIdx.x * TILE_OW;
    const int oh0 = blockIdx.y * TILE_OH;
    const int ih0 = oh0 - pad_h;
    const int iw0 = ow0 - pad_w;

    const float* in_base = input + bc * (long)in_h * in_w;

    if (ih0 >= 0 && iw0 >= 0 && ih0 + TILE_IH <= in_h && iw0 + TILE_IW <= in_w) {
        #pragma unroll
        for (int k = 0; k < (NLOAD + NT - 1) / NT; ++k) {
            const int i = tid + k * NT;
            if (i < NLOAD) {
                const int r = i / TILE_IW;
                const int col = i - r * TILE_IW;
                s_in[r][col] = in_base[(long)(ih0 + r) * in_w + (iw0 + col)];
            }
        }
    } else {
        #pragma unroll
        for (int k = 0; k < (NLOAD + NT - 1) / NT; ++k) {
            const int i = tid + k * NT;
            if (i < NLOAD) {
                const int r = i / TILE_IW;
                const int col = i - r * TILE_IW;
                const int ih = ih0 + r;
                const int iw = iw0 + col;
                float v = 0.f;
                if (ih >= 0 && ih < in_h && iw >= 0 && iw < in_w) {
                    v = in_base[(long)ih * in_w + iw];
                }
                s_in[r][col] = v;
            }
        }
    }
    __syncthreads();

    // Cache weights in registers (removes 21 shared loads per output pixel).
    float w[KH * KW];
    #pragma unroll
    for (int i = 0; i < KH * KW; ++i) w[i] = s_w[i];

    const int rb = threadIdx.y * OH;
    const int cb = threadIdx.x * OW;

    // Load the full sliding window for this thread's OW x OH outputs.
    float v[VH][W];
    #pragma unroll
    for (int i = 0; i < VH; ++i) {
        WinLoader<OW, W>::run(&s_in[rb + i][cb], v[i]);
    }

    const float b = has_bias ? bias[c] : 0.f;
    const int ow = ow0 + cb;
    const int oh = oh0 + rb;
    float* out_base = output + bc * (long)out_h * out_w;
    const bool full_tile_x = (ow + OW <= out_w);

    #pragma unroll
    for (int dy = 0; dy < OH; ++dy) {
        const int y = oh + dy;
        if (y >= out_h) continue;

        float outv[OW];
        #pragma unroll
        for (int dx = 0; dx < OW; ++dx) {
            float acc[KH];
            #pragma unroll
            for (int i = 0; i < KH; ++i) {
                acc[i] = 0.f;
                #pragma unroll
                for (int j = 0; j < KW; ++j) {
                    acc[i] += v[i + dy][j + dx] * w[i * KW + j];
                }
            }
            float sum = b;
            #pragma unroll
            for (int i = 0; i < KH; ++i) sum += acc[i];
            outv[dx] = sum;
        }

        float* out_row = out_base + (long)y * out_w + ow;
        if (full_tile_x &&
            ((reinterpret_cast<uintptr_t>(out_row) & (uintptr_t)(OW * 4 - 1)) == 0)) {
            StgV<OW>::run(out_row, outv);
        } else {
            #pragma unroll
            for (int dx = 0; dx < OW; ++dx) {
                if (ow + dx < out_w) out_row[dx] = outv[dx];
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Shared-memory tile kernel: templated on kernel size (fully unrolled),
// valid for stride == 1 and dilation == 1 (any padding). One output pixel
// per thread. Used as fallback for sizes not covered by the reg kernel.
// ---------------------------------------------------------------------------
template <int KH, int KW, typename scalar_t>
__global__ void depthwise_conv2d_fast_kernel(
    const scalar_t* __restrict__ input,
    const scalar_t* __restrict__ weight,
    const scalar_t* __restrict__ bias,
    scalar_t* __restrict__ output,
    const int channels,
    const int in_h, const int in_w,
    const int out_h, const int out_w,
    const int pad_h, const int pad_w,
    const bool has_bias)
{
    constexpr int TILE_IN_H = TILE_H + KH - 1;
    constexpr int TILE_IN_W = TILE_W + KW - 1;

    __shared__ scalar_t s_w[KH * KW];
    __shared__ scalar_t s_in[TILE_IN_H][TILE_IN_W];

    const long bc = blockIdx.x;
    const int c = (int)(bc % channels);
    const int tid = threadIdx.y * TILE_W + threadIdx.x;
    constexpr int nthreads = TILE_W * TILE_H;

    if (tid < KH * KW) {
        s_w[tid] = weight[(long)c * KH * KW + tid];
    }

    const int ow0 = blockIdx.z * TILE_W;
    const int oh0 = blockIdx.y * TILE_H;
    const int ih0 = oh0 - pad_h;
    const int iw0 = ow0 - pad_w;

    const scalar_t* in_base = input + bc * (long)in_h * in_w;

    #pragma unroll 2
    for (int i = tid; i < TILE_IN_H * TILE_IN_W; i += nthreads) {
        const int r = i / TILE_IN_W;
        const int col = i - r * TILE_IN_W;
        const int ih = ih0 + r;
        const int iw = iw0 + col;
        scalar_t v = scalar_t(0);
        if (ih >= 0 && ih < in_h && iw >= 0 && iw < in_w) {
            v = in_base[(long)ih * in_w + iw];
        }
        s_in[r][col] = v;
    }
    __syncthreads();

    const int ow = ow0 + threadIdx.x;
    const int oh = oh0 + threadIdx.y;
    if (ow < out_w && oh < out_h) {
        scalar_t sum = has_bias ? bias[c] : scalar_t(0);
        #pragma unroll
        for (int i = 0; i < KH; ++i) {
            #pragma unroll
            for (int j = 0; j < KW; ++j) {
                sum += s_in[threadIdx.y + i][threadIdx.x + j] * s_w[i * KW + j];
            }
        }
        output[(bc * (long)out_h + oh) * out_w + ow] = sum;
    }
}

// ---------------------------------------------------------------------------
// Generic kernel: arbitrary kernel size / stride / dilation / padding, using
// dynamically-sized shared memory tile.
// ---------------------------------------------------------------------------
template <typename scalar_t>
__global__ void depthwise_conv2d_generic_kernel(
    const scalar_t* __restrict__ input,
    const scalar_t* __restrict__ weight,
    const scalar_t* __restrict__ bias,
    scalar_t* __restrict__ output,
    const int channels,
    const int in_h, const int in_w,
    const int out_h, const int out_w,
    const int kh, const int kw,
    const int stride_h, const int stride_w,
    const int pad_h, const int pad_w,
    const int dil_h, const int dil_w,
    const bool has_bias)
{
    extern __shared__ char smem_raw[];
    const int tile_in_w = (TILE_W - 1) * stride_w + (kw - 1) * dil_w + 1;
    const int tile_in_h = (TILE_H - 1) * stride_h + (kh - 1) * dil_h + 1;
    scalar_t* s_w = reinterpret_cast<scalar_t*>(smem_raw);
    scalar_t* s_in = s_w + kh * kw;

    const long bc = blockIdx.x;
    const int c = (int)(bc % channels);
    const int tid = threadIdx.y * TILE_W + threadIdx.x;
    const int nthreads = TILE_W * TILE_H;

    for (int i = tid; i < kh * kw; i += nthreads) {
        s_w[i] = weight[(long)c * kh * kw + i];
    }

    const int ow0 = blockIdx.z * TILE_W;
    const int oh0 = blockIdx.y * TILE_H;
    const int ih0 = oh0 * stride_h - pad_h;
    const int iw0 = ow0 * stride_w - pad_w;

    const scalar_t* in_base = input + bc * (long)in_h * in_w;

    for (int i = tid; i < tile_in_h * tile_in_w; i += nthreads) {
        const int r = i / tile_in_w;
        const int col = i - r * tile_in_w;
        const int ih = ih0 + r;
        const int iw = iw0 + col;
        scalar_t v = scalar_t(0);
        if (ih >= 0 && ih < in_h && iw >= 0 && iw < in_w) {
            v = in_base[(long)ih * in_w + iw];
        }
        s_in[r * tile_in_w + col] = v;
    }
    __syncthreads();

    const int ow = ow0 + threadIdx.x;
    const int oh = oh0 + threadIdx.y;
    if (ow < out_w && oh < out_h) {
        scalar_t sum = has_bias ? bias[c] : scalar_t(0);
        const int r0 = threadIdx.y * stride_h;
        const int c0 = threadIdx.x * stride_w;
        for (int i = 0; i < kh; ++i) {
            const int r = r0 + i * dil_h;
            for (int j = 0; j < kw; ++j) {
                sum += s_in[r * tile_in_w + (c0 + j * dil_w)] * s_w[i * kw + j];
            }
        }
        output[(bc * (long)out_h + oh) * out_w + ow] = sum;
    }
}

// ---------------------------------------------------------------------------
// Naive fallback (used only when the shared-memory tile would be too large).
// ---------------------------------------------------------------------------
template <typename scalar_t>
__global__ void depthwise_conv2d_naive_kernel(
    const scalar_t* __restrict__ input,
    const scalar_t* __restrict__ weight,
    const scalar_t* __restrict__ bias,
    scalar_t* __restrict__ output,
    const int channels,
    const int in_h, const int in_w,
    const int out_h, const int out_w,
    const int kh, const int kw,
    const int stride_h, const int stride_w,
    const int pad_h, const int pad_w,
    const int dil_h, const int dil_w,
    const bool has_bias,
    const long total)
{
    const long idx = blockIdx.x * (long)blockDim.x + threadIdx.x;
    if (idx >= total) return;
    const int ow = (int)(idx % out_w);
    const int oh = (int)((idx / out_w) % out_h);
    const long bc = idx / ((long)out_h * out_w);
    const int c = (int)(bc % channels);

    const scalar_t* in_base = input + bc * (long)in_h * in_w;
    const scalar_t* w_base = weight + (long)c * kh * kw;

    scalar_t sum = has_bias ? bias[c] : scalar_t(0);
    for (int i = 0; i < kh; ++i) {
        const int ih = oh * stride_h - pad_h + i * dil_h;
        if (ih < 0 || ih >= in_h) continue;
        for (int j = 0; j < kw; ++j) {
            const int iw = ow * stride_w - pad_w + j * dil_w;
            if (iw < 0 || iw >= in_w) continue;
            sum += in_base[(long)ih * in_w + iw] * w_base[i * kw + j];
        }
    }
    output[idx] = sum;
}

// ---------------------------------------------------------------------------
// Register-blocked float path. Returns true when the kernel was launched.
// ---------------------------------------------------------------------------
static bool try_launch_reg_float(
    const float* input,
    const float* weight,
    const float* bias,
    float* output,
    const int batch, const int channels,
    const int in_h, const int in_w,
    const int out_h, const int out_w,
    const int kh, const int kw,
    const int pad_h, const int pad_w,
    cudaStream_t stream)
{
    const long bc_total = (long)batch * channels;
    if (bc_total > 65535) return false;
    const bool has_bias = (bias != nullptr);
    const dim3 rblock(32, 8);

#define TRY_REG(KHV, KWV, OWV, OHV)                                                \
    if (kh == KHV && kw == KWV) {                                                  \
        const dim3 rgrid(                                                          \
            (unsigned)((out_w + 32 * OWV - 1) / (32 * OWV)),                       \
            (unsigned)((out_h + 8 * OHV - 1) / (8 * OHV)),                         \
            (unsigned)bc_total);                                                   \
        if (rgrid.y > 65535) return false;                                         \
        depthwise_conv2d_reg_kernel<KHV, KWV, OWV, OHV>                            \
            <<<rgrid, rblock, 0, stream>>>(                                        \
                input, weight, bias, output, channels, in_h, in_w,                 \
                out_h, out_w, pad_h, pad_w, has_bias);                             \
        C10_CUDA_KERNEL_LAUNCH_CHECK();                                            \
        return true;                                                               \
    }

    TRY_REG(1, 3, 4, 2)
    TRY_REG(3, 1, 4, 2)
    TRY_REG(1, 5, 4, 2)
    TRY_REG(5, 1, 4, 2)
    TRY_REG(1, 7, 4, 2)
    TRY_REG(7, 1, 4, 2)
    TRY_REG(3, 3, 4, 2)
    TRY_REG(3, 5, 4, 2)
    TRY_REG(5, 3, 4, 2)
    TRY_REG(3, 7, 4, 2)
    TRY_REG(7, 3, 2, 2)
    TRY_REG(5, 5, 2, 2)
#undef TRY_REG
    return false;
}

// ---------------------------------------------------------------------------
// Launch helper
// ---------------------------------------------------------------------------
template <typename scalar_t>
static void launch_depthwise_conv2d(
    const scalar_t* input,
    const scalar_t* weight,
    const scalar_t* bias,
    scalar_t* output,
    const int batch, const int channels,
    const int in_h, const int in_w,
    const int out_h, const int out_w,
    const int kh, const int kw,
    const int stride_h, const int stride_w,
    const int pad_h, const int pad_w,
    const int dil_h, const int dil_w,
    cudaStream_t stream)
{
    const dim3 block(TILE_W, TILE_H);
    const dim3 grid(
        (unsigned int)(batch * channels),
        (unsigned int)((out_h + TILE_H - 1) / TILE_H),
        (unsigned int)((out_w + TILE_W - 1) / TILE_W));
    const bool has_bias = (bias != nullptr);

#define TRY_FAST(KHV, KWV)                                                          \
    if (kh == KHV && kw == KWV) {                                                   \
        depthwise_conv2d_fast_kernel<KHV, KWV, scalar_t>                            \
            <<<grid, block, 0, stream>>>(                                           \
                input, weight, bias, output, channels, in_h, in_w,                  \
                out_h, out_w, pad_h, pad_w, has_bias);                              \
        C10_CUDA_KERNEL_LAUNCH_CHECK();                                             \
        return;                                                                     \
    }

    if (stride_h == 1 && stride_w == 1 && dil_h == 1 && dil_w == 1) {
        TRY_FAST(1, 1)
        TRY_FAST(1, 3)
        TRY_FAST(3, 1)
        TRY_FAST(1, 5)
        TRY_FAST(5, 1)
        TRY_FAST(1, 7)
        TRY_FAST(7, 1)
        TRY_FAST(3, 3)
        TRY_FAST(3, 5)
        TRY_FAST(5, 3)
        TRY_FAST(3, 7)
        TRY_FAST(7, 3)
        TRY_FAST(5, 5)
        TRY_FAST(5, 7)
        TRY_FAST(7, 5)
        TRY_FAST(7, 7)
    }
#undef TRY_FAST

    const int tile_in_w = (TILE_W - 1) * stride_w + (kw - 1) * dil_w + 1;
    const int tile_in_h = (TILE_H - 1) * stride_h + (kh - 1) * dil_h + 1;
    const size_t shmem = (size_t)(kh * kw + tile_in_h * tile_in_w) * sizeof(scalar_t);
    if (shmem <= 48 * 1024) {
        depthwise_conv2d_generic_kernel<scalar_t><<<grid, block, shmem, stream>>>(
            input, weight, bias, output, channels, in_h, in_w, out_h, out_w,
            kh, kw, stride_h, stride_w, pad_h, pad_w, dil_h, dil_w, has_bias);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    } else {
        const long total = (long)batch * channels * out_h * out_w;
        const int threads = 256;
        const long blocks = (total + threads - 1) / threads;
        depthwise_conv2d_naive_kernel<scalar_t>
            <<<(unsigned int)blocks, threads, 0, stream>>>(
                input, weight, bias, output, channels, in_h, in_w, out_h, out_w,
                kh, kw, stride_h, stride_w, pad_h, pad_w, dil_h, dil_w,
                has_bias, total);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }
}

// ---------------------------------------------------------------------------
// Exposed op
// ---------------------------------------------------------------------------
torch::Tensor depthwise_conv2d_forward(
    const torch::Tensor& input,
    const torch::Tensor& weight,
    const c10::optional<torch::Tensor>& bias,
    const int64_t stride_h, const int64_t stride_w,
    const int64_t pad_h, const int64_t pad_w,
    const int64_t dil_h, const int64_t dil_w)
{
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
    TORCH_CHECK(input.dim() == 4, "input must be 4D (N, C, H, W)");
    TORCH_CHECK(weight.dim() == 4, "weight must be 4D (C, 1, KH, KW)");

    const c10::cuda::CUDAGuard device_guard(input.device());

    auto input_c = input.contiguous();
    auto weight_c = weight.contiguous();

    const int batch = (int)input_c.size(0);
    const int channels = (int)input_c.size(1);
    const int in_h = (int)input_c.size(2);
    const int in_w = (int)input_c.size(3);
    const int kh = (int)weight_c.size(2);
    const int kw = (int)weight_c.size(3);
    TORCH_CHECK(weight_c.size(0) == channels && weight_c.size(1) == 1,
                "expected depthwise weight of shape (C, 1, KH, KW)");

    const int out_h = (int)((in_h + 2 * pad_h - dil_h * (kh - 1) - 1) / stride_h + 1);
    const int out_w = (int)((in_w + 2 * pad_w - dil_w * (kw - 1) - 1) / stride_w + 1);
    TORCH_CHECK(out_h > 0 && out_w > 0, "invalid output size");

    auto output = torch::empty({batch, channels, out_h, out_w}, input_c.options());

    const bool has_bias = bias.has_value() && bias->defined();
    torch::Tensor bias_c;
    if (has_bias) {
        bias_c = bias->contiguous();
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    // Highest-performance path: register-blocked float kernel.
    if (input_c.scalar_type() == at::kFloat &&
        stride_h == 1 && stride_w == 1 && dil_h == 1 && dil_w == 1) {
        if (try_launch_reg_float(
                input_c.data_ptr<float>(),
                weight_c.data_ptr<float>(),
                has_bias ? bias_c.data_ptr<float>() : nullptr,
                output.data_ptr<float>(),
                batch, channels, in_h, in_w, out_h, out_w,
                kh, kw, (int)pad_h, (int)pad_w, stream)) {
            return output;
        }
    }

    AT_DISPATCH_FLOATING_TYPES(input_c.scalar_type(), "depthwise_conv2d_forward", [&] {
        launch_depthwise_conv2d<scalar_t>(
            input_c.data_ptr<scalar_t>(),
            weight_c.data_ptr<scalar_t>(),
            has_bias ? bias_c.data_ptr<scalar_t>() : nullptr,
            output.data_ptr<scalar_t>(),
            batch, channels, in_h, in_w, out_h, out_w,
            kh, kw,
            (int)stride_h, (int)stride_w,
            (int)pad_h, (int)pad_w,
            (int)dil_h, (int)dil_w,
            stream);
    });

    return output;
}
"""

_depthwise_conv2d_ext = load_inline(
    name="depthwise_conv2d_reg_ext_v3",
    cpp_sources=_depthwise_conv2d_cpp_src,
    cuda_sources=_depthwise_conv2d_cuda_src,
    functions=["depthwise_conv2d_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Depthwise 2D convolution with asymmetric input and asymmetric kernel,
    accelerated with a custom register-blocked shared-memory CUDA kernel.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size_h: int, kernel_size_w: int, stride_h: int = 1, stride_w: int = 1, padding_h: int = 0, padding_w: int = 0, dilation_h: int = 1, dilation_w: int = 1, groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, in_channels, (kernel_size_h, kernel_size_w), stride=(stride_h, stride_w), padding=(padding_h, padding_w), dilation=(dilation_h, dilation_w), groups=in_channels, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.is_cuda and (x.dtype == torch.float32 or x.dtype == torch.float64):
            return _depthwise_conv2d_ext.depthwise_conv2d_forward(
                x,
                self.conv2d.weight,
                self.conv2d.bias,
                self.conv2d.stride[0], self.conv2d.stride[1],
                self.conv2d.padding[0], self.conv2d.padding[1],
                self.conv2d.dilation[0], self.conv2d.dilation[1],
            )
        return F.conv2d(
            x,
            self.conv2d.weight,
            self.conv2d.bias,
            self.conv2d.stride,
            self.conv2d.padding,
            self.conv2d.dilation,
            self.conv2d.groups,
        )
```
Result of evaluating it:
The solution is correct on the test inputs. Mean runtime 1.93 ms; the PyTorch reference (with TF32 enabled) takes 3.7 ms; speedup 1.917x.
Profiler, top device kernels of one forward pass of your solution:
| device kernel | total device time (us) | calls |
|---|---|---|
| void depthwise_conv2d_reg_kernel<3, 7, 4, 2>(float const*, float const*, float const*, flo | 1924.1 | 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.
