```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Allow tensor-core math for any fp32 fallback matmuls.
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
# Keep fp16 GEMM reductions in fp32 for accuracy.
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

# The workload is dominated by a (128 x 32768) @ (32768 x 32768) GEMM whose fp16
# weight matrix is ~2.1 GB, i.e. it is HBM-bandwidth-bound. With batch=128 a
# single GEMM launch creates only ~256 thread blocks (128x128 tiles), barely two
# waves on a 108-SM A100, so the tail leaves the memory system idle. Splitting
# the K dimension into chunks issued as ONE strided-batched GEMM (torch.bmm ->
# cuBLAS) gives chunk-count x more thread blocks and much better SM/HBM
# utilization. cuBLAS accumulates each chunk in fp32 on tensor cores. A fused
# custom kernel then reduces the chunk partials, adds the bias, applies
# max-pool(k=2), sums over features and scales -- reading the fp16 partials once
# and writing only `batch` floats, with fp32 accumulation.
_fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cstdint>

#define THREADS 512

__device__ __forceinline__ float blockReduceSumF(float val) {
    __shared__ float shared[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_down_sync(0xffffffffu, val, offset);
    if (lane == 0) shared[wid] = val;
    __syncthreads();
    const int nwarp = (blockDim.x + 31) >> 5;
    if (wid == 0) {
        val = (lane < nwarp) ? shared[lane] : 0.0f;
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1)
            val += __shfl_down_sync(0xffffffffu, val, offset);
    }
    return val;  // only thread 0 holds the block total
}

// P: (CHUNKS, rows, N) fp16 partial sums of the K-split GEMM.
// bias: (N) fp32.
// out[r] = scale * sum_j max(bias[2j]   + sum_c P[c,r,2j],
//                            bias[2j+1] + sum_c P[c,r,2j+1])
// One float4 load of P covers 8 halfs = 4 pool pairs. Requires N % 8 == 0.
template<int CHUNKS>
__global__ void chunked_partial_maxpool2_sum_scale_kernel(
    const __half* __restrict__ P, const float* __restrict__ bias,
    float* __restrict__ out, const int rows, const long long N,
    const float scale)
{
    const int r = blockIdx.x;
    const long long row_off = (long long)r * N;
    const long long cstride = (long long)rows * N;
    const long long nv = N >> 3;  // groups of 8 halfs
    const float4* __restrict__ b4 = reinterpret_cast<const float4*>(bias);

    float acc = 0.0f;
    for (long long i = threadIdx.x; i < nv; i += blockDim.x) {
        const float4 b0 = b4[2 * i];
        const float4 b1 = b4[2 * i + 1];
        float v0 = b0.x, v1 = b0.y, v2 = b0.z, v3 = b0.w;
        float v4 = b1.x, v5 = b1.y, v6 = b1.z, v7 = b1.w;
        #pragma unroll
        for (int c = 0; c < CHUNKS; ++c) {
            const float4 raw = *reinterpret_cast<const float4*>(
                P + (long long)c * cstride + row_off + (i << 3));
            const __half2* h = reinterpret_cast<const __half2*>(&raw);
            const float2 a = __half22float2(h[0]);
            const float2 b = __half22float2(h[1]);
            const float2 d = __half22float2(h[2]);
            const float2 e = __half22float2(h[3]);
            v0 += a.x; v1 += a.y; v2 += b.x; v3 += b.y;
            v4 += d.x; v5 += d.y; v6 += e.x; v7 += e.y;
        }
        acc += fmaxf(v0, v1) + fmaxf(v2, v3)
             + fmaxf(v4, v5) + fmaxf(v6, v7);
    }
    acc = blockReduceSumF(acc);
    if (threadIdx.x == 0) out[r] = acc * scale;
}

torch::Tensor chunked_partial_maxpool2_sum_scale(torch::Tensor P,
                                                 torch::Tensor bias,
                                                 double scale) {
    TORCH_CHECK(P.is_cuda() && P.dim() == 3, "P must be a 3D CUDA tensor");
    TORCH_CHECK(P.scalar_type() == torch::kFloat16, "P must be fp16");
    TORCH_CHECK(bias.is_cuda() && bias.dim() == 1, "bias must be 1D CUDA");
    TORCH_CHECK(bias.scalar_type() == torch::kFloat32, "bias must be fp32");
    auto Pc = P.contiguous();
    auto bc = bias.contiguous();
    // float4 loads of the bias require 16B alignment.
    if (reinterpret_cast<std::uintptr_t>(bc.data_ptr<float>()) % 16 != 0) {
        bc = bc.clone();
    }
    const int chunks = Pc.size(0);
    const int rows   = Pc.size(1);
    const long long N = Pc.size(2);
    TORCH_CHECK(bc.size(0) == N, "bias size mismatch");
    TORCH_CHECK(N % 8 == 0, "N must be a multiple of 8");

    auto out = torch::empty({rows}, bc.options());
    if (rows == 0) return out;
    if (N == 0) { out.zero_(); return out; }

    const dim3 grid(rows);
    const dim3 block(THREADS);
    const float s = (float)scale;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const __half* p = reinterpret_cast<const __half*>(Pc.data_ptr<at::Half>());
    const float* bi = bc.data_ptr<float>();
    float* o = out.data_ptr<float>();

    switch (chunks) {
        case 1:
            chunked_partial_maxpool2_sum_scale_kernel<1>
                <<<grid, block, 0, stream>>>(p, bi, o, rows, N, s); break;
        case 2:
            chunked_partial_maxpool2_sum_scale_kernel<2>
                <<<grid, block, 0, stream>>>(p, bi, o, rows, N, s); break;
        case 4:
            chunked_partial_maxpool2_sum_scale_kernel<4>
                <<<grid, block, 0, stream>>>(p, bi, o, rows, N, s); break;
        case 8:
            chunked_partial_maxpool2_sum_scale_kernel<8>
                <<<grid, block, 0, stream>>>(p, bi, o, rows, N, s); break;
        case 16:
            chunked_partial_maxpool2_sum_scale_kernel<16>
                <<<grid, block, 0, stream>>>(p, bi, o, rows, N, s); break;
        default:
            TORCH_CHECK(false, "unsupported chunk count");
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_fused_cpp_source = (
    "torch::Tensor chunked_partial_maxpool2_sum_scale(torch::Tensor P, torch::Tensor bias, double scale);"
)

_fused_op = load_inline(
    name="chunked_partial_maxpool_sum_scale_v3",
    cpp_sources=_fused_cpp_source,
    cuda_sources=_fused_cuda_source,
    functions=["chunked_partial_maxpool2_sum_scale"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)

_DEFAULT_CHUNKS = 4


def _run_cublas(xh, wt, bias, scale, chunks):
    bs = xh.size(0)
    K, N = wt.shape
    if chunks == 1:
        part = torch.mm(xh, wt).unsqueeze(0)
    else:
        k = K // chunks
        # (C, bs, k) strided view; bmm passes the leading dim straight to
        # cuBLAS strided-batched GEMM -- no copies, C x more thread blocks.
        xb = xh.view(bs, chunks, k).transpose(0, 1)
        wb = wt.view(chunks, k, N)
        part = torch.bmm(xb, wb)
    return _fused_op.chunked_partial_maxpool2_sum_scale(part, bias, float(scale))


def _run_linear(xh, wh, bias, scale):
    part = F.linear(xh, wh).unsqueeze(0)
    return _fused_op.chunked_partial_maxpool2_sum_scale(part, bias, float(scale))


def _bench(fn, warmup=1, iters=3):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    s = torch.cuda.Event(enable_timing=True)
    e = torch.cuda.Event(enable_timing=True)
    s.record()
    for _ in range(iters):
        fn()
    e.record()
    torch.cuda.synchronize()
    return s.elapsed_time(e) / iters


# Global plan cache so repeated ModelNew instances with the same shapes skip
# re-tuning. Key: (device, batch, out_features, in_features).
_PLAN_CACHE = {}


def _select_plan(owner, xh, bias):
    M = xh.size(0)
    K, N = owner._wT_half.shape
    key = (xh.device.index, M, N, K)
    plan = _PLAN_CACHE.get(key)
    if plan is not None:
        return plan

    wt = owner._wT_half
    scale = float(owner.scale_factor)
    cands = []
    for c in (1, 2, 4, 8, 16):
        if K % c == 0:
            cands.append(("cublas", c))
    cands.append(("linear", None))

    # Known-good default if timing fails for every candidate.
    best_plan = ("cublas", _DEFAULT_CHUNKS if K % _DEFAULT_CHUNKS == 0 else 1)
    best_t = float("inf")
    for kind, cfg in cands:
        try:
            if kind == "cublas":
                fn = lambda c=cfg: _run_cublas(xh, wt, bias, scale, c)
            else:
                fn = lambda: _run_linear(xh, owner._get_wh(), bias, scale)
            t = _bench(fn)
            if t < best_t:
                best_t = t
                best_plan = (kind, cfg)
        except Exception:
            continue
    _PLAN_CACHE[key] = best_plan
    return best_plan


class ModelNew(nn.Module):
    """
    Optimized model:
      - The linear layer runs in fp16 on tensor cores. The workload is bound by
        streaming the ~2.1 GB fp16 weight from HBM, so several cuBLAS execution
        plans are benchmarked once (cached globally per shape) and the fastest
        wins:
          * K-split strided-batched GEMM (torch.bmm, chunk counts 1/2/4/8/16)
            whose fp16 partials feed a fused custom CUDA kernel doing
            partial-reduction + bias-add + max-pool(k=2) + feature-sum + scale
            with fp32 accumulation.
          * single TN GEMM via F.linear + the same fused epilogue.
      - The transposed fp16 weight (chunk-major layout) is built once and
        cached, keyed on the parameter version and device.
      - Non-CUDA / unsupported configurations fall back to the original ops.
    """
    def __init__(self, in_features, out_features, kernel_size, scale_factor):
        super(ModelNew, self).__init__()
        self.matmul = nn.Linear(in_features, out_features)
        self.kernel_size = kernel_size
        self.scale_factor = scale_factor
        self._wT_half = None
        self._w_half = None
        self._w_ver = -1

    def _refresh_cache(self, w):
        ver = w._version
        if (self._wT_half is None or self._w_ver != ver
                or self._wT_half.device != w.device):
            # (in_f, out_f) fp16, contiguous (NN layout); viewing as
            # (C, in_f/C, out_f) yields per-chunk weight^T slices for free.
            self._wT_half = w.detach().half().t().contiguous()
            self._w_half = None
            self._w_ver = ver

    def _get_wh(self):
        w = self.matmul.weight
        if self._w_half is None:
            # (out_f, in_f) fp16 (TN layout, just a cast of the original).
            self._w_half = w.detach().half().contiguous()
        return self._w_half

    def forward(self, x):
        w = self.matmul.weight
        b = self.matmul.bias
        n = w.size(0)
        if (x.is_cuda and x.dim() == 2 and x.size(0) > 0
                and x.dtype == torch.float32 and b is not None
                and self.kernel_size == 2 and n % 8 == 0):
            try:
                self._refresh_cache(w)
                xh = x.half()
                if not xh.is_contiguous():
                    xh = xh.contiguous()
                bias = b.detach()
                if not bias.is_contiguous():
                    bias = bias.contiguous()
                kind, cfg = _select_plan(self, xh, bias)
                scale = float(self.scale_factor)
                if kind == "cublas":
                    return _run_cublas(xh, self._wT_half, bias, scale, cfg)
                return _run_linear(xh, self._get_wh(), bias, scale)
            except Exception:
                pass
        # Fallback (CPU / other dtypes / unsupported shapes): original ops.
        y = F.linear(x, w, b)
        y = F.max_pool1d(y.unsqueeze(1), self.kernel_size).squeeze(1)
        y = torch.sum(y, dim=1)
        return y * self.scale_factor
```