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 3 of 6 for this task. Your previous solution (turn 2) 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>
#include <type_traits>

#define TILE_W 32
#define TILE_H 16

// ---------------------------------------------------------------------------
// Vectorized shared-memory load / global-store helpers (float only).
// Pointers must satisfy the natural alignment of the widest chunk used.
// ---------------------------------------------------------------------------
template <int N>
__device__ __forceinline__ void lds_vec4(const float* __restrict__ p, float* __restrict__ v) {
    if constexpr (N >= 4) {
        *reinterpret_cast<float4*>(v) = *reinterpret_cast<const float4*>(p);
        lds_vec4<N - 4>(p + 4, v + 4);
    } else if constexpr (N >= 2) {
        *reinterpret_cast<float2*>(v) = *reinterpret_cast<const float2*>(p);
        lds_vec4<N - 2>(p + 2, v + 2);
    } else if constexpr (N == 1) {
        v[0] = p[0];
    }
}

template <int N>
__device__ __forceinline__ void lds_vec2(const float* __restrict__ p, float* __restrict__ v) {
    if constexpr (N >= 2) {
        *reinterpret_cast<float2*>(v) = *reinterpret_cast<const float2*>(p);
        lds_vec2<N - 2>(p + 2, v + 2);
    } else if constexpr (N == 1) {
        v[0] = p[0];
    }
}

template <int N>
__device__ __forceinline__ void stg_vec4(float* __restrict__ p, const float* __restrict__ v) {
    if constexpr (N >= 4) {
        *reinterpret_cast<float4*>(p) = *reinterpret_cast<const float4*>(v);
        stg_vec4<N - 4>(p + 4, v + 4);
    } else if constexpr (N >= 2) {
        *reinterpret_cast<float2*>(p) = *reinterpret_cast<const float2*>(v);
        stg_vec4<N - 2>(p + 2, v + 2);
    } else if constexpr (N == 1) {
        p[0] = v[0];
    }
}

template <int N>
__device__ __forceinline__ void stg_vec2(float* __restrict__ p, const float* __restrict__ v) {
    if constexpr (N >= 2) {
        *reinterpret_cast<float2*>(p) = *reinterpret_cast<const float2*>(v);
        stg_vec2<N - 2>(p + 2, v + 2);
    } else if constexpr (N == 1) {
        p[0] = v[0];
    }
}

// ---------------------------------------------------------------------------
// 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 input window is
// loaded once into registers with wide shared-memory loads (sliding window
// reuse across all outputs of the thread), which drastically cuts 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 to multiple of 4 floats
    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).
    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) {
        if constexpr (OW == 4) {
            lds_vec4<W>(&s_in[rb + i][cb], v[i]);
        } else {
            lds_vec2<W>(&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) {
            const uintptr_t a = reinterpret_cast<uintptr_t>(out_row);
            if ((a & 15) == 0) {
                stg_vec4<OW>(out_row, outv);
            } else if ((a & 7) == 0) {
                stg_vec2<OW>(out_row, outv);
            } else {
                #pragma unroll
                for (int dx = 0; dx < OW; ++dx) out_row[dx] = outv[dx];
            }
        } 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;
}

// ---------------------------------------------------------------------------
// 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);
    const long bc_total = (long)batch * channels;

    // Highest-performance path: register-blocked float kernel.
    if constexpr (std::is_same<scalar_t, float>::value) {
        if (stride_h == 1 && stride_w == 1 && dil_h == 1 && dil_w == 1 &&
            bc_total <= 65535) {
            const unsigned gy = (unsigned)((out_h + 15) / 16);
            if (gy <= 65535) {
                const dim3 rblock(32, 8);
#define TRY_REG(KHV, KWV, OWV)                                                     \
    if (kh == KHV && kw == KWV) {                                                  \
        const dim3 rgrid(                                                          \
            (unsigned)((out_w + 32 * OWV - 1) / (32 * OWV)),                       \
            gy,                                                                    \
            (unsigned)bc_total);                                                   \
        depthwise_conv2d_reg_kernel<KHV, KWV, OWV, 2>                              \
            <<<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;                                                                    \
    }
                TRY_REG(1, 3, 4)
                TRY_REG(3, 1, 4)
                TRY_REG(1, 5, 4)
                TRY_REG(5, 1, 4)
                TRY_REG(1, 7, 4)
                TRY_REG(7, 1, 4)
                TRY_REG(3, 3, 4)
                TRY_REG(3, 5, 4)
                TRY_REG(5, 3, 4)
                TRY_REG(3, 7, 4)
                TRY_REG(7, 3, 2)
                TRY_REG(5, 5, 2)
#undef TRY_REG
            }
        }
    }

#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();

    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_tiled_ext_v2",
    cpp_sources=_depthwise_conv2d_cpp_src,
    cuda_sources=_depthwise_conv2d_cuda_src,
    functions=["depthwise_conv2d_forward"],
    verbose=False,
    extra_cflags=["-O3", "-std=c++17"],
    extra_cuda_cflags=["-O3", "-std=c++17"],
)


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 failed to compile. Error output:
```
...[truncated]
 ?const c10::intrusive_ptr<torch::autograd::Node>? to ?const c10::SymInt&?
  493 | DECLARE_SYMINT_OP(float, SymFloat) // just for completeness
      | ^~~~~~~~~~~~~~~~~
/opt/kbv/v/lib/python3.12/site-packages/torch/include/c10/core/SymInt.h:493:1: note: candidate: ?bool c10::operator!=(float, const SymInt&)?
  493 | DECLARE_SYMINT_OP(float, SymFloat) // just for completeness
      | ^~~~~~~~~~~~~~~~~
/opt/kbv/v/lib/python3.12/site-packages/torch/include/c10/core/SymInt.h:493:1: note:   no known conversion for argument 2 from ?std::nullptr_t? to ?const c10::SymInt&?
  493 | DECLARE_SYMINT_OP(float, SymFloat) // just for completeness
      | ^~~~~~~~~~~~~~~~~
/opt/kbv/v/lib/python3.12/site-packages/torch/include/c10/core/ScalarTypeToTypeMeta.h:50:13: note: candidate: ?bool c10::operator!=(ScalarType, caffe2::TypeMeta)?
   50 | inline bool operator!=(ScalarType t, caffe2::TypeMeta m) {
      |             ^~~~~~~~
/opt/kbv/v/lib/python3.12/site-packages/torch/include/c10/core/ScalarTypeToTypeMeta.h:50:35: note:   no known conversion for argument 1 from ?const c10::intrusive_ptr<torch::autograd::Node>? to ?c10::ScalarType?
   50 | inline bool operator!=(ScalarType t, caffe2::TypeMeta m) {
      |                        ~~~~~~~~~~~^
/opt/kbv/v/lib/python3.12/site-packages/torch/include/c10/core/ScalarTypeToTypeMeta.h:54:13: note: candidate: ?bool c10::operator!=(caffe2::TypeMeta, ScalarType)?
   54 | inline bool operator!=(caffe2::TypeMeta m, ScalarType t) {
      |             ^~~~~~~~
/opt/kbv/v/lib/python3.12/site-packages/torch/include/c10/core/ScalarTypeToTypeMeta.h:54:41: note:   no known conversion for argument 1 from ?const c10::intrusive_ptr<torch::autograd::Node>? to ?caffe2::TypeMeta?
   54 | inline bool operator!=(caffe2::TypeMeta m, ScalarType t) {
      |                        ~~~~~~~~~~~~~~~~~^
/opt/kbv/v/lib/python3.12/site-packages/torch/include/c10/core/Allocator.h:152:13: note: candidate: ?bool c10::operator!=(const DataPtr&, std::nullptr_t)?
  152 | inline bool operator!=(const DataPtr& dp, std::nullptr_t) noexcept {
      |             ^~~~~~~~
/opt/kbv/v/lib/python3.12/site-packages/torch/include/c10/core/Allocator.h:152:39: note:   no known conversion for argument 1 from ?const c10::intrusive_ptr<torch::autograd::Node>? to ?const c10::DataPtr&?
  152 | inline bool operator!=(const DataPtr& dp, std::nullptr_t) noexcept {
      |                        ~~~~~~~~~~~~~~~^~
/opt/kbv/v/lib/python3.12/site-packages/torch/include/c10/core/Allocator.h:155:13: note: candidate: ?bool c10::operator!=(std::nullptr_t, const DataPtr&)?
  155 | inline bool operator!=(std::nullptr_t, const DataPtr& dp) noexcept {
      |             ^~~~~~~~
/opt/kbv/v/lib/python3.12/site-packages/torch/include/c10/core/Allocator.h:155:24: note:   no known conversion for argument 1 from ?const c10::intrusive_ptr<torch::autograd::Node>? to ?std::nullptr_t?
  155 | inline bool operator!=(std::nullptr_t, const DataPtr& dp) noexcept {
      |                        ^~~~~~~~~~~~~~
In file included from /opt/kbv/v/lib/python3.12/site-packages/torch/include/torch/csrc/autograd/input_metadata.h:4,
                 from /opt/kbv/v/lib/python3.12/site-packages/torch/include/torch/csrc/autograd/node.h:8,
                 from /opt/kbv/v/lib/python3.12/site-packages/torch/include/torch/csrc/autograd/variable.h:11:
/opt/kbv/v/lib/python3.12/site-packages/torch/include/ATen/ExpandUtils.h: In function ?at::Tensor at::_sum_to(Tensor, c10::ArrayRef<T>, bool)?:
/opt/kbv/v/lib/python3.12/site-packages/torch/include/ATen/ExpandUtils.h:470:28: error: ?sizes? is not a member of ?at::symint?
  470 |   auto sizes = at::symint::sizes<T>(tensor);
      |                            ^~~~~
/opt/kbv/v/lib/python3.12/site-packages/torch/include/ATen/ExpandUtils.h:470:35: error: expected primary-expression before ?>? token
  470 |   auto sizes = at::symint::sizes<T>(tensor);
      |                                   ^
ninja: build stopped: subcommand failed.

```

The fastest correct solution so far (1.630x speedup) is:
```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>

#define TILE_W 32
#define TILE_H 16

// ---------------------------------------------------------------------------
// Fast path kernel: templated on kernel size (fully unrolled), valid for
// stride == 1 and dilation == 1 (any padding). Uses a shared-memory input
// tile plus shared-memory weights. Each thread produces one output pixel.
// Grid: (batch*channels, ceil(out_h/TILE_H), ceil(out_w/TILE_W))
// Block: (TILE_W, TILE_H)
// ---------------------------------------------------------------------------
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;
}

// ---------------------------------------------------------------------------
// 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();

    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_tiled_ext",
    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 tiled 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,
        )
```

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.
