```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 split-K matrix-vector multiplication kernel (C = A @ B).
# A is (M, K) row-major, B is (K, 1), C is (M, 1).
#
# The operation is purely DRAM-bandwidth bound (A is streamed exactly once),
# so the kernel is engineered for maximum sustainable read throughput:
#   * 8 rows per thread block: the (L2-resident) B slice is loaded once and
#     reused for all 8 rows, cutting B load instructions and L2 traffic by
#     8x versus one row per block while keeping 8 independent 128-bit A
#     loads in flight per thread.
#   * The grid is ordered split-K-part-fastest (blockIdx.x = K slice,
#     blockIdx.y = row group) so that the ~few hundred concurrently resident
#     blocks densely cover a small, contiguous slab of A (a few full row
#     groups) instead of scattering thousands of streams across the whole
#     multi-GB matrix. This greatly improves DRAM page locality and TLB
#     reach, approaching memcpy-class streaming bandwidth.
#   * A is read with 128-bit __ldcs (evict-first streaming) loads so it
#     does not displace the heavily reused B in L2.
#   * The split-K factor is chosen at runtime from the measured occupancy
#     and SM count to minimize tail-wave quantization; partials are
#     combined with atomicAdd into a zero-initialized output.
_matvec_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <climits>

#define MV_THREADS 256
#define MV_ROWS 8

__global__ void __launch_bounds__(MV_THREADS) matvec_rows8_kernel(
    const float4* __restrict__ A4,
    const float4* __restrict__ B4,
    const float*  __restrict__ As,
    const float*  __restrict__ Bs,
    float* __restrict__ C,
    const int K4,
    const int K,
    const int chunk,
    const int M)
{
    const int row0  = blockIdx.y * MV_ROWS;
    const int rem   = M - row0;              // 1..MV_ROWS valid rows
    const int part  = blockIdx.x;
    const int start = part * chunk;
    int end = start + chunk;
    if (end > K4) end = K4;

    const float4* Ab = A4 + (size_t)row0 * (size_t)K4;

    // Per-row float4 offsets; extra rows clamp to the last valid row and
    // their results are discarded (no atomic issued for them).
    int off[MV_ROWS];
    float acc[MV_ROWS];
    #pragma unroll
    for (int r = 0; r < MV_ROWS; ++r) {
        const int rr = (r < rem) ? r : (rem - 1);
        off[r] = rr * K4;
        acc[r] = 0.f;
    }

    // Main streaming loop: one cached B float4 feeds 8 streaming A loads.
    for (int i = start + threadIdx.x; i < end; i += MV_THREADS) {
        const float4 b = __ldg(B4 + i);
        #pragma unroll
        for (int r = 0; r < MV_ROWS; ++r) {
            const float4 a = __ldcs(Ab + (off[r] + i));
            acc[r] += a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
        }
    }

    // Scalar tail when K is not divisible by 4 (handled by slice 0).
    if (part == 0) {
        const int base = K4 * 4;
        if (base < K) {
            for (int j = base + threadIdx.x; j < K; j += MV_THREADS) {
                const float bj = __ldg(Bs + j);
                #pragma unroll
                for (int r = 0; r < MV_ROWS; ++r) {
                    const int rr = (r < rem) ? r : (rem - 1);
                    acc[r] += __ldg(As + (size_t)(row0 + rr) * (size_t)K + j) * bj;
                }
            }
        }
    }

    // Warp-level reduction of all 8 row sums.
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        #pragma unroll
        for (int r = 0; r < MV_ROWS; ++r)
            acc[r] += __shfl_down_sync(0xffffffffu, acc[r], o);
    }

    __shared__ float w[MV_THREADS / 32][MV_ROWS];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane == 0) {
        #pragma unroll
        for (int r = 0; r < MV_ROWS; ++r) w[warp][r] = acc[r];
    }
    __syncthreads();

    if (threadIdx.x < MV_ROWS) {
        const int r = threadIdx.x;
        float v = 0.f;
        #pragma unroll
        for (int t = 0; t < MV_THREADS / 32; ++t) v += w[t][r];
        if (r < rem) atomicAdd(C + row0 + r, v);
    }
}

static int g_blocks_per_sm = 0;

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");
    TORCH_CHECK(A.dim() == 2, "A must be 2-D");
    TORCH_CHECK(B.numel() == A.size(1), "B must have A.size(1) elements");

    const c10::cuda::CUDAGuard device_guard(A.device());

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

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

    if (M == 0) return torch::empty({M, 1}, Ac.options());
    if (K == 0) return torch::zeros({M, 1}, Ac.options());

    const int64_t row_groups = (M + MV_ROWS - 1) / MV_ROWS;
    const int64_t K4 = K >> 2;

    // The vectorized kernel requires 16-byte-aligned operands (true for
    // fresh allocations), 32-bit in-block index arithmetic, and a grid.y
    // within limits; otherwise fall back to cuBLAS.
    if ((reinterpret_cast<uintptr_t>(Ac.data_ptr<float>()) & 0xF) != 0 ||
        (reinterpret_cast<uintptr_t>(Bc.data_ptr<float>()) & 0xF) != 0 ||
        K4 * MV_ROWS >= (int64_t)INT32_MAX ||
        row_groups > 65535 || M > (int64_t)INT32_MAX ||
        K > (int64_t)INT32_MAX) {
        return torch::matmul(A, B);
    }

    auto C = torch::empty({M, 1}, Ac.options());
    auto stream = at::cuda::getCurrentCUDAStream();
    cudaMemsetAsync(C.data_ptr<float>(), 0, sizeof(float) * (size_t)M,
                    stream.stream());

    if (g_blocks_per_sm == 0) {
        int n = 0;
        cudaOccupancyMaxActiveBlocksPerMultiprocessor(
            &n, matvec_rows8_kernel, MV_THREADS, 0);
        g_blocks_per_sm = (n > 0) ? n : 4;
    }
    const int sm_count =
        at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    const int64_t cap = (int64_t)g_blocks_per_sm * sm_count;

    // Pick the split-K factor that minimizes tail-wave quantization: model
    // the makespan as ceil(blocks / resident_capacity) * chunk and take the
    // best candidate (ties favor the smaller split = fewer blocks).
    int split = 1, chunk = 0;
    if (K4 > 0) {
        long long best_eff = LLONG_MAX;
        for (int s = 1; s <= 128; s <<= 1) {
            int ch = (int)((K4 + s - 1) / s);
            ch = ((ch + MV_THREADS - 1) / MV_THREADS) * MV_THREADS;
            const int sp = (int)((K4 + ch - 1) / ch);
            const long long blocks = (long long)row_groups * sp;
            const long long waves = (blocks + cap - 1) / cap;
            const long long eff = waves * (long long)ch;
            if (eff < best_eff) {
                best_eff = eff;
                split = sp;
                chunk = ch;
            }
        }
    }

    dim3 grid((unsigned)split, (unsigned)row_groups);
    dim3 block(MV_THREADS);

    matvec_rows8_kernel<<<grid, block, 0, stream>>>(
        reinterpret_cast<const float4*>(Ac.data_ptr<float>()),
        reinterpret_cast<const float4*>(Bc.data_ptr<float>()),
        Ac.data_ptr<float>(),
        Bc.data_ptr<float>(),
        C.data_ptr<float>(),
        (int)K4, (int)K, chunk, (int)M);

    return C;
}
"""

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

_matvec_ext = load_inline(
    name="matvec_splitk_ext_v6",
    cpp_sources=_matvec_cpp_src,
    cuda_sources=_matvec_cuda_src,
    functions=["matvec_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Matrix-vector multiplication (C = A * B) accelerated with a custom
    split-K CUDA kernel: 8 rows per block, split-K-part-fastest grid order
    for dense contiguous DRAM streaming of A, cached loads for the
    L2-resident B, and an occupancy-aware split-K factor that minimizes
    tail-wave quantization.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._op = _matvec_ext

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if (A.is_cuda and B.is_cuda and A.dim() == 2
                and A.dtype == torch.float32 and B.dtype == torch.float32):
            return self._op.matvec_cuda(A, B)
        return torch.matmul(A, B)
```