```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 (GEMV): C = A @ B
# A: (M, K), B: (K, 1), C: (M, 1)
# This operation is heavily memory-bound (A is read exactly once), so the
# kernel is optimized for maximum DRAM bandwidth utilization:
#   - one thread block per output row (M = 2048 blocks -> full GPU occupancy)
#   - vectorized 16-byte float4 loads for both A and B
#   - 4x loop unrolling to keep multiple loads in flight (memory-level parallelism)
#   - B stays resident in L2 cache and is reused across all row-blocks
#   - warp-shuffle based block reduction (no bank conflicts, no __syncthreads in loop)

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

#define BLOCK_SIZE 256
#define NUM_WARPS (BLOCK_SIZE / 32)

__device__ __forceinline__ float warp_reduce_sum(float v) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        v += __shfl_xor_sync(0xffffffffu, v, off);
    }
    return v;
}

__device__ __forceinline__ float block_reduce_sum(float v, float* smem) {
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    v = warp_reduce_sum(v);
    if (lane == 0) smem[wid] = v;
    __syncthreads();
    if (wid == 0) {
        float s = (lane < NUM_WARPS) ? smem[lane] : 0.0f;
        s = warp_reduce_sum(s);
        return s;  // valid on lane 0
    }
    return 0.0f;
}

// Vectorized kernel: requires K % 4 == 0 (16-byte aligned rows)
__global__ void gemv_vec4_kernel(const float4* __restrict__ A,
                                 const float4* __restrict__ B,
                                 float* __restrict__ C,
                                 const long K4) {
    const float4* __restrict__ a = A + (long)blockIdx.x * K4;
    float sum = 0.0f;

    const long stride = blockDim.x;
    long i = threadIdx.x;

    // 4x unrolled main loop for memory-level parallelism
    const long limit = K4 - 3 * stride;
    for (; i < limit; i += 4 * stride) {
        float4 a0 = a[i];
        float4 b0 = B[i];
        float4 a1 = a[i + stride];
        float4 b1 = B[i + stride];
        float4 a2 = a[i + 2 * stride];
        float4 b2 = B[i + 2 * stride];
        float4 a3 = a[i + 3 * stride];
        float4 b3 = B[i + 3 * stride];
        sum += a0.x * b0.x + a0.y * b0.y + a0.z * b0.z + a0.w * b0.w;
        sum += a1.x * b1.x + a1.y * b1.y + a1.z * b1.z + a1.w * b1.w;
        sum += a2.x * b2.x + a2.y * b2.y + a2.z * b2.z + a2.w * b2.w;
        sum += a3.x * b3.x + a3.y * b3.y + a3.z * b3.z + a3.w * b3.w;
    }
    for (; i < K4; i += stride) {
        float4 av = a[i];
        float4 bv = B[i];
        sum += av.x * bv.x + av.y * bv.y + av.z * bv.z + av.w * bv.w;
    }

    __shared__ float smem[NUM_WARPS];
    float total = block_reduce_sum(sum, smem);
    if (threadIdx.x == 0) {
        C[blockIdx.x] = total;
    }
}

// Scalar fallback kernel for arbitrary K
__global__ void gemv_scalar_kernel(const float* __restrict__ A,
                                   const float* __restrict__ B,
                                   float* __restrict__ C,
                                   const long K) {
    const float* __restrict__ a = A + (long)blockIdx.x * K;
    float sum = 0.0f;
    for (long i = threadIdx.x; i < K; i += blockDim.x) {
        sum += a[i] * B[i];
    }

    __shared__ float smem[NUM_WARPS];
    float total = block_reduce_sum(sum, smem);
    if (threadIdx.x == 0) {
        C[blockIdx.x] = total;
    }
}

torch::Tensor gemv_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
                "Only float32 is supported");

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

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

    auto C = torch::empty({M, 1}, Ac.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if ((K & 3) == 0) {
        const long K4 = K >> 2;
        gemv_vec4_kernel<<<(int)M, BLOCK_SIZE, 0, stream>>>(
            reinterpret_cast<const float4*>(Ac.data_ptr<float>()),
            reinterpret_cast<const float4*>(Bc.data_ptr<float>()),
            C.data_ptr<float>(),
            K4);
    } else {
        gemv_scalar_kernel<<<(int)M, BLOCK_SIZE, 0, stream>>>(
            Ac.data_ptr<float>(),
            Bc.data_ptr<float>(),
            C.data_ptr<float>(),
            K);
    }

    return C;
}
"""

gemv_cpp_source = "torch::Tensor gemv_cuda(torch::Tensor A, torch::Tensor B);"

# Compile the inline CUDA GEMV extension
gemv_ext = load_inline(
    name="gemv_ext",
    cpp_sources=gemv_cpp_source,
    cuda_sources=gemv_cuda_source,
    functions=["gemv_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs matrix-vector multiplication (C = A * B)
    using a custom bandwidth-optimized CUDA GEMV kernel.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.gemv = gemv_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
                and A.dim() == 2):
            return self.gemv.gemv_cuda(A, B)
        # Fallback for non-CUDA / non-float32 inputs
        return torch.matmul(A, B)
```