```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

# Custom CUDA kernel for matrix-vector multiplication (C = A @ B).
# The operation is memory-bandwidth bound, so the kernel is designed to:
#   1. Use one thread block per output row (M = 2048 blocks -> good occupancy).
#   2. Use vectorized float4 loads for both A and B (128-bit transactions).
#   3. Unroll by 4 with independent accumulators for memory-level parallelism.
#   4. Use warp-shuffle + shared-memory block reduction (no atomics).
#   5. Use 64-bit indexing since M*K exceeds INT32 range.
matvec_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

#define BLOCK_SIZE 256

__device__ __forceinline__ float warp_reduce(float v) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        v += __shfl_down_sync(0xffffffffu, v, offset);
    }
    return v;
}

__device__ __forceinline__ float block_reduce(float v) {
    __shared__ float shared[32];
    int lane = threadIdx.x & 31;
    int wid  = threadIdx.x >> 5;

    v = warp_reduce(v);
    if (lane == 0) shared[wid] = v;
    __syncthreads();

    int nw = (blockDim.x + 31) >> 5;
    v = (threadIdx.x < nw) ? shared[threadIdx.x] : 0.0f;
    if (wid == 0) {
        v = warp_reduce(v);
    }
    return v;  // result valid in thread 0
}

// Vectorized path (requires K % 4 == 0 and contiguous inputs).
__global__ void matvec_vec4_kernel(const float* __restrict__ A,
                                   const float* __restrict__ B,
                                   float* __restrict__ C,
                                   int M, int64_t K) {
    int row = blockIdx.x;
    if (row >= M) return;

    const float* a_row = A + (int64_t)row * K;
    const float4* a4 = reinterpret_cast<const float4*>(a_row);
    const float4* b4 = reinterpret_cast<const float4*>(B);
    const int64_t K4 = K >> 2;

    float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f;
    const int stride = BLOCK_SIZE;
    int64_t i = threadIdx.x;
    const int64_t limit = K4 - 3 * (int64_t)stride;

    // Unrolled main loop: 4 independent float4 loads per thread per iter.
    for (; i < limit; i += 4 * (int64_t)stride) {
        float4 a0 = a4[i];
        float4 a1 = a4[i + stride];
        float4 a2 = a4[i + 2 * stride];
        float4 a3 = a4[i + 3 * stride];
        float4 b0 = b4[i];
        float4 b1 = b4[i + stride];
        float4 b2 = b4[i + 2 * stride];
        float4 b3 = b4[i + 3 * stride];
        s0 += a0.x * b0.x + a0.y * b0.y + a0.z * b0.z + a0.w * b0.w;
        s1 += a1.x * b1.x + a1.y * b1.y + a1.z * b1.z + a1.w * b1.w;
        s2 += a2.x * b2.x + a2.y * b2.y + a2.z * b2.z + a2.w * b2.w;
        s3 += a3.x * b3.x + a3.y * b3.y + a3.z * b3.z + a3.w * b3.w;
    }
    // Remaining float4 chunks.
    for (; i < K4; i += stride) {
        float4 a = a4[i];
        float4 b = b4[i];
        s0 += a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
    }
    // Scalar tail (in case K is not a multiple of 4).
    for (int64_t j = (K4 << 2) + threadIdx.x; j < K; j += stride) {
        s1 += a_row[j] * B[j];
    }

    float sum = (s0 + s1) + (s2 + s3);
    sum = block_reduce(sum);
    if (threadIdx.x == 0) C[row] = sum;
}

// Fully general scalar fallback path.
__global__ void matvec_scalar_kernel(const float* __restrict__ A,
                                     const float* __restrict__ B,
                                     float* __restrict__ C,
                                     int M, int64_t K) {
    int row = blockIdx.x;
    if (row >= M) return;

    const float* a_row = A + (int64_t)row * K;
    float sum = 0.0f;
    for (int64_t i = threadIdx.x; i < K; i += blockDim.x) {
        sum += a_row[i] * B[i];
    }

    sum = block_reduce(sum);
    if (threadIdx.x == 0) C[row] = sum;
}

torch::Tensor matvec_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "A and B must be CUDA tensors");
    TORCH_CHECK(A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
                "A and B must be float32");

    auto Ac = A.contiguous();
    auto Bc = B.contiguous();

    const int64_t M = Ac.size(0);
    const int64_t K = Ac.size(1);

    auto C = torch::empty({M, 1}, Ac.options());

    dim3 block(BLOCK_SIZE);
    dim3 grid((unsigned int)M);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if ((K & 3) == 0) {
        matvec_vec4_kernel<<<grid, block, 0, stream>>>(
            Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(),
            (int)M, K);
    } else {
        matvec_scalar_kernel<<<grid, block, 0, stream>>>(
            Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(),
            (int)M, K);
    }

    return C;
}
"""

matvec_cpp_source = "torch::Tensor matvec_cuda(torch::Tensor A, torch::Tensor B);"

# Compile the inline CUDA extension.
matvec_ext = load_inline(
    name="matvec_custom_ext",
    cpp_sources=matvec_cpp_source,
    cuda_sources=matvec_cuda_source,
    functions=["matvec_cuda"],
    verbose=True,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs matrix-vector multiplication (C = A * B)
    using a custom bandwidth-optimized CUDA kernel.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.matvec = matvec_ext

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if (A.is_cuda and B.is_cuda
                and A.dtype == torch.float32 and B.dtype == torch.float32):
            return self.matvec.matvec_cuda(A, B)
        # Fallback for non-CUDA or non-float32 inputs.
        return torch.matmul(A, B)
```