```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline


_cpp_source = r"""
#include <torch/extension.h>

torch::Tensor asymmetric_conv_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias);
"""


_cuda_source = 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 <mma.h>

namespace wmma = nvcuda::wmma;

// Convert contiguous NCHW input to NHWC with coalesced reads and writes.
__global__ void pack_input(
    const float* __restrict__ x,
    float* __restrict__ packed,
    int spatial)
{
    __shared__ float tile[32][33];

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int p0 = blockIdx.x * 32;
    const int batch = blockIdx.y;

    #pragma unroll
    for (int j = 0; j < 32; j += 8) {
        const int p = p0 + tx;
        tile[ty + j][tx] =
            p < spatial
                ? x[((int64_t)batch * 32 + ty + j) * spatial + p]
                : 0.0f;
    }
    __syncthreads();

    #pragma unroll
    for (int j = 0; j < 32; j += 8) {
        const int p = p0 + ty + j;
        if (p < spatial) {
            packed[((int64_t)batch * spatial + p) * 32 + tx] =
                tile[tx][ty + j];
        }
    }
}

// OIHW -> (kernel_height, kernel_width, input_channel, output_channel).
__global__ void pack_weight(
    const float* __restrict__ weight,
    float* __restrict__ packed)
{
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < 45 * 32 * 64) {
        const int oc = i % 64;
        const int ic = (i / 64) % 32;
        const int filter = i / (64 * 32);
        packed[i] = weight[(oc * 32 + ic) * 45 + filter];
    }
}

// Each block computes 64 spatial positions by 64 output channels.
// Three TF32 products retain substantially more precision than one
// TF32 product: Ah*Bh + Ah*Bl + Al*Bh.
__global__ void conv_5x9_32x64(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int height,
    int width,
    int out_height,
    int out_width)
{
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    __shared__ __align__(32) float storage[8192];

    float* ah = storage;
    float* al = storage + 2048;
    float* bh = storage + 4096;
    float* bl = storage + 6144;

    const int tid = threadIdx.x;
    const int warp = tid / 32;
    const int tile_row = (warp / 4) * 32;
    const int tile_col = (warp % 4) * 16;

    const int batch = blockIdx.y;
    const int position_base = blockIdx.x * 64;
    const int out_spatial = out_height * out_width;

    using AFragment = wmma::fragment<
        wmma::matrix_a, 16, 16, 8,
        wmma::precision::tf32, wmma::row_major>;
    using BFragment = wmma::fragment<
        wmma::matrix_b, 16, 16, 8,
        wmma::precision::tf32, wmma::row_major>;
    using AccFragment = wmma::fragment<
        wmma::accumulator, 16, 16, 8, float>;

    AccFragment main_acc[2];
    AccFragment correction[2];

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        wmma::fill_fragment(main_acc[i], 0.0f);
        wmma::fill_fragment(correction[i], 0.0f);
    }

    #pragma unroll 1
    for (int filter = 0; filter < 45; ++filter) {
        const int fh = filter / 9;
        const int fw = filter % 9;

        #pragma unroll
        for (int i = tid; i < 2048; i += 256) {
            const int position = position_base + i / 32;
            const int channel = i % 32;
            float value = 0.0f;

            if (position < out_spatial) {
                const int row = position / out_width;
                const int col = position % out_width;
                const int64_t offset =
                    (((int64_t)batch * height + row + fh) * width
                     + col + fw) * 32 + channel;
                value = x[offset];
            }

            const float high = __float_to_tf32(value);
            ah[i] = high;
            al[i] = __float_to_tf32(value - high);

            const float w = weight[filter * 2048 + i];
            const float w_high = __float_to_tf32(w);
            bh[i] = w_high;
            bl[i] = __float_to_tf32(w - w_high);
        }

        __syncthreads();

        #pragma unroll
        for (int k = 0; k < 32; k += 8) {
            BFragment b_high;
            BFragment b_low;
            wmma::load_matrix_sync(
                b_high, bh + k * 64 + tile_col, 64);
            wmma::load_matrix_sync(
                b_low, bl + k * 64 + tile_col, 64);

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                AFragment a_high;
                AFragment a_low;
                const int offset = (tile_row + i * 16) * 32 + k;

                wmma::load_matrix_sync(a_high, ah + offset, 32);
                wmma::load_matrix_sync(a_low, al + offset, 32);

                wmma::mma_sync(
                    main_acc[i], a_high, b_high, main_acc[i]);
                wmma::mma_sync(
                    correction[i], a_high, b_low, correction[i]);
                wmma::mma_sync(
                    correction[i], a_low, b_high, correction[i]);
            }
        }

        __syncthreads();
    }

    // Reuse shared memory for the output transpose.
    constexpr int output_stride = 68;

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < AccFragment::num_elements; ++j) {
            main_acc[i].x[j] += correction[i].x[j];
        }

        wmma::store_matrix_sync(
            storage + (tile_row + i * 16) * output_stride + tile_col,
            main_acc[i],
            output_stride,
            wmma::mem_row_major);
    }

    __syncthreads();

    #pragma unroll
    for (int i = tid; i < 4096; i += 256) {
        const int oc = i / 64;
        const int local_position = i % 64;
        const int position = position_base + local_position;

        if (position < out_spatial) {
            float value = storage[local_position * output_stride + oc];
            if (bias != nullptr) {
                value += bias[oc];
            }
            output[((int64_t)batch * 64 + oc) * out_spatial + position] =
                value;
        }
    }
#endif
}

torch::Tensor asymmetric_conv_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "Expected CUDA tensors");
    TORCH_CHECK(
        x.scalar_type() == at::kFloat &&
        weight.scalar_type() == at::kFloat,
        "Expected float32 tensors");
    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous(),
                "Expected contiguous tensors");
    TORCH_CHECK(x.dim() == 4 && x.size(1) == 32,
                "Expected NCHW input with 32 channels");
    TORCH_CHECK(
        weight.dim() == 4 && weight.size(0) == 64 &&
        weight.size(1) == 32 && weight.size(2) == 5 &&
        weight.size(3) == 9,
        "Expected weight shape (64, 32, 5, 9)");
    TORCH_CHECK(x.device() == weight.device(), "Device mismatch");
    TORCH_CHECK(x.size(2) >= 5 && x.size(3) >= 9,
                "Input is smaller than the kernel");

    const c10::cuda::CUDAGuard guard(x.device());

    const int batch = static_cast<int>(x.size(0));
    const int height = static_cast<int>(x.size(2));
    const int width = static_cast<int>(x.size(3));
    const int out_height = height - 4;
    const int out_width = width - 8;
    const int spatial = height * width;
    const int out_spatial = out_height * out_width;

    auto output = torch::empty(
        {batch, 64, out_height, out_width}, x.options());

    if (batch == 0) {
        return output;
    }

    const float* bias_ptr = nullptr;
    if (bias.has_value() && bias->defined()) {
        TORCH_CHECK(
            bias->device() == x.device() &&
            bias->scalar_type() == at::kFloat &&
            bias->is_contiguous() && bias->numel() == 64,
            "Invalid bias");
        bias_ptr = bias->data_ptr<float>();
    }

    auto packed_input = torch::empty(
        {batch, height, width, 32}, x.options());
    auto packed_weight = torch::empty(
        {45, 32, 64}, weight.options());

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

    pack_input<<<
        dim3((spatial + 31) / 32, batch), dim3(32, 8), 0, stream>>>(
        x.data_ptr<float>(),
        packed_input.data_ptr<float>(),
        spatial);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    pack_weight<<<(45 * 32 * 64 + 255) / 256, 256, 0, stream>>>(
        weight.data_ptr<float>(),
        packed_weight.data_ptr<float>());
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    conv_5x9_32x64<<<
        dim3((out_spatial + 63) / 64, batch), 256, 0, stream>>>(
        packed_input.data_ptr<float>(),
        packed_weight.data_ptr<float>(),
        bias_ptr,
        output.data_ptr<float>(),
        height,
        width,
        out_height,
        out_width);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""


_extension = None


def _get_extension():
    global _extension
    if _extension is None:
        _extension = load_inline(
            name="asymmetric_conv_tf32x3_v1",
            cpp_sources=_cpp_source,
            cuda_sources=_cuda_source,
            functions=["asymmetric_conv_cuda"],
            extra_cflags=["-O3"],
            extra_cuda_cflags=["-O3"],
            verbose=False,
        )
    return _extension


class ModelNew(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: tuple,
        stride: int = 1,
        padding: int = 0,
        dilation: int = 1,
        groups: int = 1,
        bias: bool = False,
    ):
        super().__init__()
        self.conv2d = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=bias,
        )
        conv = self.conv2d
        self._specialized = (
            in_channels == 32
            and out_channels == 64
            and conv.kernel_size == (5, 9)
            and conv.stride == (1, 1)
            and conv.padding == (0, 0)
            and conv.dilation == (1, 1)
            and groups == 1
        )
        self._device_support = {}

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        weight = self.conv2d.weight

        if (
            not self._specialized
            or not x.is_cuda
            or x.dtype != torch.float32
            or weight.dtype != torch.float32
            or x.device != weight.device
            or x.ndim != 4
            or not x.is_contiguous()
            or not weight.is_contiguous()
            or torch.is_autocast_enabled()
        ):
            return self.conv2d(x)

        device = x.device.index
        if device not in self._device_support:
            self._device_support[device] = (
                torch.cuda.get_device_capability(device)[0] >= 8
            )

        if not self._device_support[device]:
            return self.conv2d(x)

        return _get_extension().asymmetric_conv_cuda(
            x, weight, self.conv2d.bias
        )
```