```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 TF32 tensor-core math for FP32 matmuls (explicitly permitted by the
# precision rule). Storage stays FP32.
torch.backends.cuda.matmul.allow_tf32 = True
try:
    torch.set_float32_matmul_precision("high")
except Exception:
    pass

# ---------------------------------------------------------------------------
# Custom kernels:
#  1) gemm_tf32_splitk_kernel: TF32 tensor-core (wmma) GEMM specialized for
#     M=128, computing partial C = A @ W^T over a K-chunk with split-K so the
#     whole GPU is occupied (the default cuBLAS pick leaves ~40% of the SMs
#     idle for this shape). The weight matrix is streamed from HBM with
#     cp.async double buffering.
#  2) reduce_min_sub_kernel: deterministic split-K reduction fused with the
#     epilogue: out = min(sum(partials) + bias, c) - c in a single pass.
# ---------------------------------------------------------------------------
gemm_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <mma.h>
#include <algorithm>
#include <cmath>

using namespace nvcuda;

#define BM 128
#define BN 128
#define BK 64
#define STAGES 2
#define LDS (BK + 4)
#define NTHREADS 256
#define SMEM_BYTES (STAGES * (BM + BN) * LDS * 4)

__device__ __forceinline__ void cp_async16(void* smem_ptr, const void* gmem_ptr) {
    unsigned saddr = static_cast<unsigned>(__cvta_generic_to_shared(smem_ptr));
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: "r"(saddr), "l"(gmem_ptr));
}
__device__ __forceinline__ void cp_async_commit() {
    asm volatile("cp.async.commit_group;\n");
}
template <int N>
__device__ __forceinline__ void cp_async_wait() {
    asm volatile("cp.async.wait_group %0;\n" :: "n"(N));
}

// C_partial = A[128,K] @ W[N,K]^T for a K-chunk, written to ws[s][128][N].
__global__ void __launch_bounds__(NTHREADS)
gemm_tf32_splitk_kernel(const float* __restrict__ A,
                        const float* __restrict__ W,
                        float* __restrict__ ws,
                        const int gridN, const int K, const int N,
                        const int base_iters, const int rem) {
    const int s  = blockIdx.x / gridN;
    const int nt = blockIdx.x % gridN;
    const int n0 = nt * BN;

    const int kiters = base_iters + (s < rem ? 1 : 0);
    const int kstart = s * base_iters + min(s, rem);

    extern __shared__ float smem[];
    float* As[STAGES];
    float* Ws[STAGES];
    for (int st = 0; st < STAGES; ++st) {
        As[st] = smem + st * (BM + BN) * LDS;
        Ws[st] = As[st] + BM * LDS;
    }

    const int tid = threadIdx.x;
    const int warp = tid >> 5;
    const int warpM = warp >> 2;   // 0..1  (64 rows each)
    const int warpN = warp & 3;    // 0..3  (32 cols each)

    // Accumulators: 4 (M) x 2 (N) 16x16 tiles per warp.
    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc[4][2];
    for (int mi = 0; mi < 4; ++mi)
        for (int ni = 0; ni < 2; ++ni)
            wmma::fill_fragment(acc[mi][ni], 0.0f);

    const long Kstride = K;

    auto prefetch = [&](int iter, int buf) {
        const int k0 = (kstart + iter) * BK;
        // A tile: rows 0..127, 16 float4 chunks per row
        for (int t = tid; t < BM * (BK / 4); t += NTHREADS) {
            const int row = t >> 4;
            const int ch  = t & 15;
            cp_async16(As[buf] + row * LDS + ch * 4,
                       A + (long)row * Kstride + k0 + ch * 4);
        }
        // W tile: rows n0..n0+127
        for (int t = tid; t < BN * (BK / 4); t += NTHREADS) {
            const int row = t >> 4;
            const int ch  = t & 15;
            cp_async16(Ws[buf] + row * LDS + ch * 4,
                       W + (long)(n0 + row) * Kstride + k0 + ch * 4);
        }
        cp_async_commit();
    };

    prefetch(0, 0);

    for (int i = 0; i < kiters; ++i) {
        const int buf = i & (STAGES - 1);
        if (i + 1 < kiters) {
            prefetch(i + 1, (i + 1) & (STAGES - 1));
            cp_async_wait<1>();
        } else {
            cp_async_wait<0>();
        }
        __syncthreads();

        const float* a = As[buf];
        const float* w = Ws[buf];
        #pragma unroll
        for (int kk = 0; kk < BK / 8; ++kk) {
            wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> fa[4];
            wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::col_major> fb[2];
            #pragma unroll
            for (int mi = 0; mi < 4; ++mi) {
                wmma::load_matrix_sync(fa[mi], a + (warpM * 64 + mi * 16) * LDS + kk * 8, LDS);
                #pragma unroll
                for (int e = 0; e < fa[mi].num_elements; ++e)
                    fa[mi].x[e] = wmma::__float_to_tf32(fa[mi].x[e]);
            }
            #pragma unroll
            for (int ni = 0; ni < 2; ++ni) {
                wmma::load_matrix_sync(fb[ni], w + (warpN * 32 + ni * 16) * LDS + kk * 8, LDS);
                #pragma unroll
                for (int e = 0; e < fb[ni].num_elements; ++e)
                    fb[ni].x[e] = wmma::__float_to_tf32(fb[ni].x[e]);
            }
            #pragma unroll
            for (int mi = 0; mi < 4; ++mi)
                #pragma unroll
                for (int ni = 0; ni < 2; ++ni)
                    wmma::mma_sync(acc[mi][ni], fa[mi], fb[ni], acc[mi][ni]);
        }
        __syncthreads();
    }

    // Stage the 128x128 accumulator tile through shared memory for fully
    // coalesced global writes.
    float* smC = smem;  // reuse (>= 128*132 floats)
    #pragma unroll
    for (int mi = 0; mi < 4; ++mi)
        #pragma unroll
        for (int ni = 0; ni < 2; ++ni)
            wmma::store_matrix_sync(
                smC + (warpM * 64 + mi * 16) * 132 + warpN * 32 + ni * 16,
                acc[mi][ni], 132, wmma::mem_row_major);
    __syncthreads();

    float* out_base = ws + ((long)s * BM) * N + n0;
    for (int t = tid; t < BM * (BN / 4); t += NTHREADS) {
        const int row = t >> 5;
        const int c4  = t & 31;
        *reinterpret_cast<float4*>(out_base + (long)row * N + c4 * 4) =
            *reinterpret_cast<const float4*>(smC + row * 132 + c4 * 4);
    }
}

// out = min(sum_s ws[s] + bias, c) - c   (vectorized float4)
__global__ void reduce_min_sub_kernel(const float4* __restrict__ ws,
                                      const float4* __restrict__ bias,
                                      float4* __restrict__ out,
                                      const int S,
                                      const long n4,
                                      const long biasN4,
                                      const float c) {
    const long stride = (long)gridDim.x * blockDim.x;
    for (long i = blockIdx.x * (long)blockDim.x + threadIdx.x; i < n4; i += stride) {
        float4 acc = bias[i % biasN4];
        const float4* p = ws + i;
        for (int sp = 0; sp < S; ++sp) {
            float4 v = p[(long)sp * n4];
            acc.x += v.x; acc.y += v.y; acc.z += v.z; acc.w += v.w;
        }
        acc.x = fminf(acc.x, c) - c;
        acc.y = fminf(acc.y, c) - c;
        acc.z = fminf(acc.z, c) - c;
        acc.w = fminf(acc.w, c) - c;
        out[i] = acc;
    }
}

torch::Tensor gemm_min_sub_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, double c) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda() && b.is_cuda(), "tensors must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");

    const long M = x.size(0);
    const long K = x.size(1);
    const long N = w.size(0);
    TORCH_CHECK(M == BM, "M must be 128");
    TORCH_CHECK(K % BK == 0, "K must be a multiple of 64");
    TORCH_CHECK(N % BN == 0, "N must be a multiple of 128");

    auto xc = x.contiguous();
    auto wc = w.contiguous();
    auto bc = b.contiguous();

    const int gridN = (int)(N / BN);
    const int total_iters = (int)(K / BK);

    int dev = 0;
    cudaGetDevice(&dev);
    int sms = 0;
    cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
    if (sms <= 0) sms = 108;

    // Pick the split-K factor that minimizes wave-quantization waste.
    int bestS = 1;
    double bestEff = -1.0;
    const int maxS = std::min(16, total_iters);
    for (int s = 1; s <= maxS; ++s) {
        const long blocks = (long)gridN * s;
        const long waves = (blocks + sms - 1) / sms;
        const double eff = (double)blocks / (double)(waves * sms);
        if (eff > bestEff + 1e-9) { bestEff = eff; bestS = s; }
    }
    const int S = bestS;
    const int base_iters = total_iters / S;
    const int rem = total_iters % S;

    auto ws = at::empty({(long)S, M, N}, x.options());
    auto out = at::empty({M, N}, x.options());

    static bool attr_set = false;
    if (!attr_set) {
        cudaFuncSetAttribute(gemm_tf32_splitk_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             SMEM_BYTES);
        attr_set = true;
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    dim3 grid((unsigned)((long)gridN * S));
    gemm_tf32_splitk_kernel<<<grid, NTHREADS, SMEM_BYTES, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), ws.data_ptr<float>(),
        gridN, (int)K, (int)N, base_iters, rem);

    const long n4 = (M * N) >> 2;
    const long biasN4 = N >> 2;
    const long blocks_needed = (n4 + NTHREADS - 1) / NTHREADS;
    const int rgrid = (int)std::min<long>(blocks_needed, 8192L);
    reduce_min_sub_kernel<<<rgrid, NTHREADS, 0, stream>>>(
        reinterpret_cast<const float4*>(ws.data_ptr<float>()),
        reinterpret_cast<const float4*>(bc.data_ptr<float>()),
        reinterpret_cast<float4*>(out.data_ptr<float>()),
        S, n4, biasN4, static_cast<float>(c));

    return out;
}
"""

gemm_cpp_source = (
    "torch::Tensor gemm_min_sub_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, double c);"
)

fused_gemm = load_inline(
    name="fused_gemm_tf32_splitk_min_sub",
    cpp_sources=gemm_cpp_source,
    cuda_sources=gemm_source,
    functions=["gemm_min_sub_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    - GEMM (x @ W^T) runs on a custom TF32 tensor-core kernel with split-K
      chosen to fill every SM (the cuBLAS pick leaves SMs idle at this shape)
      and cp.async double-buffered weight streaming.
    - Bias add, torch.min(x, constant) and x - constant are all fused into
      the split-K reduction kernel, so the elementwise epilogue is free.
    - Falls back to the PyTorch path for any shape the kernel does not
      specialize.
    """
    def __init__(self, in_features, out_features, constant):
        super(ModelNew, self).__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.constant = nn.Parameter(torch.tensor(constant))
        self._const_val = float(constant)

    def forward(self, x):
        if (x.is_cuda and x.dtype == torch.float32
                and x.size(0) == 128 and x.dim() == 2
                and self.linear.in_features % 64 == 0
                and self.linear.out_features % 128 == 0):
            return fused_gemm.gemm_min_sub_cuda(
                x, self.linear.weight, self.linear.bias, self._const_val)
        out = F.linear(x, self.linear.weight, self.linear.bias)
        return torch.min(out, self.constant) - self.constant
```