```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

# Fused matrix-vector multiplication (C = A @ B) CUDA kernels.
#
# Workload: A (M, K) = (2048, 1048576), B (K, 1) -> C (M, 1).
# Purely DRAM-bandwidth bound (~8 GiB of A must be streamed exactly once),
# so the whole design is about keeping the memory controllers saturated:
#   * float4 (128-bit) fully-coalesced loads with deep unrolling for
#     memory-level parallelism.
#   * __ldcs (evict-first streaming) for A: every element is read once and
#     must not displace B in the cache hierarchy.
#   * __ldg (read-only, L1-cached) for B: with fine K-slices, each block's
#     B slice is small enough to live in L1/L2 and be shared by all blocks
#     scheduled in the same slice column.
#   * Split-K structure: K is partitioned into many slices, giving thousands
#     of small, equal-sized work items. Fine slices (16-64) keep tail-wave
#     quantization losses low and load balance near-perfect even when the
#     grid is many waves deep; each block atomically accumulates its partial
#     sums into C (one tiny async memset beforehand).
#   * Block sizes / rows-per-block / unroll / slice-count are all templated;
#     the best combination is chosen once per process by a short autotuning
#     pass at module construction time on proxy tensors.

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

namespace {

// ------------------------------ split-K kernel ------------------------------
// grid = (ceil(M / R), num_slices); each block computes partial dot products
// for R rows over one slice of K, then atomically accumulates into C.
template <int BLOCK, int R, int UNROLL>
__global__ void __launch_bounds__(BLOCK)
matvec_splitk_kernel(const float4* __restrict__ A4,
                     const float4* __restrict__ B4,
                     float* __restrict__ C,
                     const int M,
                     const int K4,
                     const int slice4) {
    constexpr int NW = BLOCK / 32;
    const int row0 = blockIdx.x * R;
    const int i0   = blockIdx.y * slice4;
    const int i1   = min(K4, i0 + slice4);

    const float4* a_rows[R];
    #pragma unroll
    for (int r = 0; r < R; ++r) {
        int row = row0 + r;
        if (row >= M) row = M - 1;  // tail: valid reads, masked on write
        a_rows[r] = A4 + static_cast<size_t>(row) * K4;
    }

    float4 acc[R];
    #pragma unroll
    for (int r = 0; r < R; ++r) acc[r] = make_float4(0.f, 0.f, 0.f, 0.f);

    const int stride = BLOCK;
    int i = i0 + threadIdx.x;

    for (; i + (UNROLL - 1) * stride < i1; i += UNROLL * stride) {
        float4 b[UNROLL];
        #pragma unroll
        for (int u = 0; u < UNROLL; ++u) b[u] = __ldg(B4 + i + u * stride);
        #pragma unroll
        for (int u = 0; u < UNROLL; ++u) {
            #pragma unroll
            for (int r = 0; r < R; ++r) {
                const float4 a = __ldcs(a_rows[r] + i + u * stride);
                acc[r].x += a.x * b[u].x;
                acc[r].y += a.y * b[u].y;
                acc[r].z += a.z * b[u].z;
                acc[r].w += a.w * b[u].w;
            }
        }
    }
    for (; i < i1; i += stride) {
        const float4 b = __ldg(B4 + i);
        #pragma unroll
        for (int r = 0; r < R; ++r) {
            const float4 a = __ldcs(a_rows[r] + i);
            acc[r].x += a.x * b.x;
            acc[r].y += a.y * b.y;
            acc[r].z += a.z * b.z;
            acc[r].w += a.w * b.w;
        }
    }

    float sum[R];
    #pragma unroll
    for (int r = 0; r < R; ++r) sum[r] = (acc[r].x + acc[r].y) + (acc[r].z + acc[r].w);

    #pragma unroll
    for (int r = 0; r < R; ++r) {
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            sum[r] += __shfl_down_sync(0xffffffffu, sum[r], off);
    }

    __shared__ float ws[R][NW];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) {
        #pragma unroll
        for (int r = 0; r < R; ++r) ws[r][wid] = sum[r];
    }
    __syncthreads();

    if (wid == 0) {
        #pragma unroll
        for (int r = 0; r < R; ++r) {
            float s = (lane < NW) ? ws[r][lane] : 0.f;
            #pragma unroll
            for (int off = NW / 2; off > 0; off >>= 1)
                s += __shfl_down_sync(0xffffffffu, s, off);
            if (lane == 0) {
                const int row = row0 + r;
                if (row < M) atomicAdd(&C[row], s);
            }
        }
    }
}

template <int BLOCK, int R, int UNROLL>
void launch_splitk(const float4* A4, const float4* B4, float* C,
                   int64_t M, int64_t K, int split_target,
                   cudaStream_t stream) {
    const int K4 = static_cast<int>(K >> 2);
    int slice4 = (K4 + split_target - 1) / split_target;
    const int align = BLOCK * UNROLL;
    slice4 = ((slice4 + align - 1) / align) * align;
    const int num_slices = (K4 + slice4 - 1) / slice4;

    const dim3 grid(static_cast<unsigned int>((M + R - 1) / R),
                    static_cast<unsigned int>(num_slices));
    C10_CUDA_CHECK(cudaMemsetAsync(C, 0, static_cast<size_t>(M) * sizeof(float), stream));
    matvec_splitk_kernel<BLOCK, R, UNROLL><<<grid, BLOCK, 0, stream>>>(
        A4, B4, C, static_cast<int>(M), K4, slice4);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
}

}  // namespace

int64_t matvec_num_configs() { return 12; }

torch::Tensor matvec_cuda(torch::Tensor A, torch::Tensor B, int64_t config) {
    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");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2 && B.size(1) == 1,
                "Expected shapes A (M, K) and B (K, 1)");
    TORCH_CHECK(A.size(1) == B.size(0), "Inner dimensions must match");

    const int64_t M = A.size(0);
    const int64_t K = A.size(1);
    auto C = torch::empty({M, 1}, A.options());

    if (M == 0 || K == 0) {
        return C.zero_();
    }
    if ((K & 3) != 0) {
        // Alignment fallback (never triggered by this model's shapes).
        C.copy_(torch::matmul(A, B));
        return C;
    }

    const float4* A4 = reinterpret_cast<const float4*>(A.data_ptr<float>());
    const float4* B4 = reinterpret_cast<const float4*>(B.data_ptr<float>());
    float* Cp = C.data_ptr<float>();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    switch (config) {
        case 0:  launch_splitk<256, 4, 4>(A4, B4, Cp, M, K, 32, stream); break;
        case 1:  launch_splitk<256, 4, 8>(A4, B4, Cp, M, K, 32, stream); break;
        case 2:  launch_splitk<512, 4, 4>(A4, B4, Cp, M, K, 32, stream); break;
        case 3:  launch_splitk<512, 4, 2>(A4, B4, Cp, M, K, 32, stream); break;
        case 4:  launch_splitk<256, 4, 4>(A4, B4, Cp, M, K, 16, stream); break;
        case 5:  launch_splitk<512, 4, 4>(A4, B4, Cp, M, K, 16, stream); break;
        case 6:  launch_splitk<512, 4, 2>(A4, B4, Cp, M, K, 16, stream); break;
        case 7:  launch_splitk<256, 2, 8>(A4, B4, Cp, M, K, 32, stream); break;
        case 8:  launch_splitk<256, 8, 4>(A4, B4, Cp, M, K, 16, stream); break;
        case 9:  launch_splitk<128, 4, 8>(A4, B4, Cp, M, K, 32, stream); break;
        case 11: launch_splitk<256, 4, 4>(A4, B4, Cp, M, K, 64, stream); break;
        case 10:
        default: launch_splitk<512, 4, 2>(A4, B4, Cp, M, K, 8, stream);  break;
    }
    return C;
}
"""

matvec_cpp_source = (
    "torch::Tensor matvec_cuda(torch::Tensor A, torch::Tensor B, int64_t config);\n"
    "int64_t matvec_num_configs();"
)

matvec_ext = load_inline(
    name="matvec_stream_ext_v3",
    cpp_sources=matvec_cpp_source,
    cuda_sources=matvec_cuda_source,
    functions=["matvec_cuda", "matvec_num_configs"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)

_K_PROXY = 1048576  # keep the real K so per-block streaming behavior matches
_TUNED_CONFIG = None


def _tune_on_proxy(proxy_m):
    """Benchmark every kernel config on a proxy problem; return fastest id."""
    dev = torch.device("cuda", torch.cuda.current_device())
    A = torch.rand(proxy_m, _K_PROXY, device=dev)
    B = torch.rand(_K_PROXY, 1, device=dev)
    ref = None
    try:
        ref = torch.matmul(A, B)
        tol = 1e-2 * max(ref.abs().max().item(), 1.0)
        n = matvec_ext.matvec_num_configs()
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        best_cfg, best_t = 10, None
        for c in range(n):
            try:
                out = matvec_ext.matvec_cuda(A, B, c)
                torch.cuda.synchronize()
                err = (out - ref).abs().max().item()
                del out
                if err > tol:
                    continue  # numerically broken config; never use it
                for _ in range(2):
                    matvec_ext.matvec_cuda(A, B, c)
                torch.cuda.synchronize()
                start.record()
                for _ in range(3):
                    matvec_ext.matvec_cuda(A, B, c)
                end.record()
                torch.cuda.synchronize()
                t = start.elapsed_time(end)
                if best_t is None or t < best_t:
                    best_t, best_cfg = t, c
            except Exception:
                continue
        return best_cfg
    finally:
        del A, B, ref
        torch.cuda.empty_cache()


def _autotune():
    global _TUNED_CONFIG
    if _TUNED_CONFIG is not None:
        return _TUNED_CONFIG
    cfg = 10
    if torch.cuda.is_available():
        for proxy_m in (2048, 1024, 256):
            try:
                cfg = _tune_on_proxy(proxy_m)
                break
            except Exception:
                torch.cuda.empty_cache()
                cfg = 10
    _TUNED_CONFIG = cfg
    return cfg


class ModelNew(nn.Module):
    """
    Optimized matrix-vector multiplication (C = A @ B) backed by a custom
    bandwidth-saturating split-K CUDA kernel (geometry autotuned once at
    construction). Falls back to torch.matmul for any unsupported input
    configuration.
    """

    def __init__(self):
        super(ModelNew, self).__init__()
        self._fn = matvec_ext.matvec_cuda
        self._cfg = _autotune()

    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
            and B.dim() == 2
            and B.size(1) == 1
            and A.size(1) == B.size(0)
        ):
            return self._fn(A.contiguous(), B.contiguous(), self._cfg)
        return torch.matmul(A, B)
```