```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 and compute type stay FP32.
torch.backends.cuda.matmul.allow_tf32 = True
try:
    torch.set_float32_matmul_precision("high")
except Exception:
    pass

# ---------------------------------------------------------------------------
# Fully fused implementation of:  out = min(x @ W.T + bias, c) - c
#
# For the benchmark shape (M=128, K=N=16384) the op is dominated by streaming
# the 1 GiB FP32 weight matrix from HBM exactly once. This kernel:
#   * TF32 tensor cores via raw mma.sync.m16n8k8 PTX (no slow wmma smem path),
#     accumulator fragments are kept resident in registers across the whole
#     K-chunk (zero fragment store/reload traffic).
#   * 5-stage cp.async pipeline with an L2 "sliding window": the 8 MiB
#     activation matrix is marked evict_last so it stays in L2 and is re-served
#     to all blocks, while streamed weights use the normal policy.
#   * Persistent grid + split-K: 1024 (n-tile, k-chunk) work items are
#     distributed over a grid sized to the SM count, so there is no wave
#     quantization tail; a grid-stride loop also removes per-iteration index
#     arithmetic from the hot loop.
#   * A runtime cost model picks the split-K factor S; for the target shape it
#     picks S=8. Partials are written fully coalesced (float4, via a smem
#     staging tile), and a fused reduction applies bias/min/subtract while
#     streaming the partials at HBM speed.
# ---------------------------------------------------------------------------
gemm_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <algorithm>

#define BM 128
#define BN 128
#define BK 32
#define STAGES 5
#define LDS 40
#define NTHREADS 256
#define STAGE_FLOATS ((BM + BN) * LDS)
#define SMEM_BYTES (STAGES * STAGE_FLOATS * 4)
#define CPLDS (LDS / 4)

__device__ __forceinline__ void cp_async16(void* smem_ptr, const void* gmem_ptr,
                                           unsigned long long policy) {
    unsigned saddr = static_cast<unsigned>(__cvta_generic_to_shared(smem_ptr));
    asm volatile("cp.async.cg.shared.global.L2::cache_hint [%0], [%1], 16, %2;\n"
                 :: "r"(saddr), "l"(gmem_ptr), "l"(policy));
}
__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));
}

__device__ __forceinline__ void mma_tf32(float& d0, float& d1, float& d2, float& d3,
                                         float a0, float a1, float a2, float a3,
                                         float b0, float b1) {
    asm volatile(
        "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
        : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
        : "r"(__float_as_uint(a0)), "r"(__float_as_uint(a1)),
          "r"(__float_as_uint(a2)), "r"(__float_as_uint(a3)),
          "r"(__float_as_uint(b0)), "r"(__float_as_uint(b1)));
}

// Each block handles a contiguous range of work items; each item computes a
// 128x128 output tile over a K/S chunk and writes raw partials to ws[s].
__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 kiters_per, const int total_items) {
    unsigned long long pol_a;
    asm volatile("createpolicy.fractional.L2::evict_last.b64 %0, 1.0;\n" : "=l"(pol_a));

    extern __shared__ float smem[];

    const int tid  = threadIdx.x;
    const int warp = tid >> 5;
    const int lane = tid & 31;
    const int g  = lane >> 2;   // fragment row group 0..7
    const int tg = lane & 3;    // fragment col group 0..3
    const int warpM = warp >> 2;   // 0..1 -> 64 rows each
    const int warpN = warp & 3;    // 0..3 -> 32 cols each

    const int stage_loads = (BM + BN) * (BK / 4);   // float4 chunks per stage
    const long Kstride = K;

    for (int item = blockIdx.x; item < total_items; item += gridDim.x) {
        const int s  = item / gridN;
        const int nt = item - s * gridN;
        const int n0 = nt * BN;
        const int k_iter0 = s * kiters_per;

        float acc[4][4][4];
        #pragma unroll
        for (int mi = 0; mi < 4; ++mi)
            #pragma unroll
            for (int ni = 0; ni < 4; ++ni)
                #pragma unroll
                for (int e = 0; e < 4; ++e) acc[mi][ni][e] = 0.0f;

        // Prologue: fill STAGES-1 pipeline stages.
        #pragma unroll
        for (int st = 0; st < STAGES - 1; ++st) {
            const int k0 = (k_iter0 + st) * BK;
            float* smA = smem + st * STAGE_FLOATS;
            float* smW = smA + BM * LDS;
            #pragma unroll
            for (int t = tid; t < stage_loads; t += NTHREADS) {
                const int row = t / CPLDS;
                const int ch  = t - row * CPLDS;
                if (row < BM)
                    cp_async16(smA + row * LDS + ch * 4, A + (long)row * Kstride + k0 + ch * 4, pol_a);
                else
                    cp_async16(smW + (row - BM) * LDS + ch * 4, W + (long)(n0 + (row - BM)) * Kstride + k0 + ch * 4, pol_a);
            }
            cp_async_commit();
        }

        // Mainloop in blocks of STAGES so the buffer index is compile-time.
        const int iters = kiters_per;
        int i = 0;
        for (; i + STAGES <= iters; i += STAGES) {
            #pragma unroll
            for (int st = 0; st < STAGES; ++st) {
                cp_async_wait<STAGES - 2>();
                __syncthreads();

                const float* smA = smem + st * STAGE_FLOATS;
                const float* smW = smA + BM * LDS;
                const float* a_base = smA + warpM * 64 * LDS + g * LDS + tg * 2;
                const float* b_base = smW + warpN * 32 * LDS + tg * 2 + g * LDS;

                float fa[2][4][4];
                float fb[2][4][2];

                #pragma unroll
                for (int mi = 0; mi < 4; ++mi) {
                    const float* p = a_base + mi * 16 * LDS;
                    float2 v01 = *reinterpret_cast<const float2*>(p);
                    float2 v23 = *reinterpret_cast<const float2*>(p + 8 * LDS);
                    fa[0][mi][0] = v01.x; fa[0][mi][1] = v23.x;
                    fa[0][mi][2] = v01.y; fa[0][mi][3] = v23.y;
                }
                #pragma unroll
                for (int ni = 0; ni < 4; ++ni) {
                    float2 v01 = *reinterpret_cast<const float2*>(b_base + ni * 16 * LDS);
                    fb[0][ni][0] = v01.x; fb[0][ni][1] = v01.y;
                }

                #pragma unroll
                for (int kk = 0; kk < BK / 8; ++kk) {
                    if (kk < BK / 8 - 1) {
                        #pragma unroll
                        for (int mi = 0; mi < 4; ++mi) {
                            const float* p = a_base + (kk + 1) * 8 + mi * 16 * LDS;
                            float2 v01 = *reinterpret_cast<const float2*>(p);
                            float2 v23 = *reinterpret_cast<const float2*>(p + 8 * LDS);
                            fa[(kk + 1) & 1][mi][0] = v01.x; fa[(kk + 1) & 1][mi][1] = v23.x;
                            fa[(kk + 1) & 1][mi][2] = v01.y; fa[(kk + 1) & 1][mi][3] = v23.y;
                        }
                        #pragma unroll
                        for (int ni = 0; ni < 4; ++ni) {
                            float2 v01 = *reinterpret_cast<const float2*>(b_base + (kk + 1) * 8 + ni * 16 * LDS);
                            fb[(kk + 1) & 1][ni][0] = v01.x; fb[(kk + 1) & 1][ni][1] = v01.y;
                        }
                    }
                    const int buf = kk & 1;
                    #pragma unroll
                    for (int mi = 0; mi < 4; ++mi)
                        #pragma unroll
                        for (int ni = 0; ni < 4; ++ni)
                            mma_tf32(acc[mi][ni][0], acc[mi][ni][1], acc[mi][ni][2], acc[mi][ni][3],
                                     fa[buf][mi][0], fa[buf][mi][1], fa[buf][mi][2], fa[buf][mi][3],
                                     fb[buf][ni][0], fb[buf][ni][1]);
                }

                // Prefetch stage i + STAGES - 1 into buffer st - 1 (mod STAGES).
                const int nxt = i + st + STAGES - 1;
                if (nxt < iters) {
                    const int k0 = (k_iter0 + nxt) * BK;
                    const int pb = (st == 0) ? (STAGES - 1) : (st - 1);
                    float* pA = smem + pb * STAGE_FLOATS;
                    float* pW = pA + BM * LDS;
                    #pragma unroll
                    for (int t = tid; t < stage_loads; t += NTHREADS) {
                        const int row = t / CPLDS;
                        const int ch  = t - row * CPLDS;
                        if (row < BM)
                            cp_async16(pA + row * LDS + ch * 4, A + (long)row * Kstride + k0 + ch * 4, pol_a);
                        else
                            cp_async16(pW + (row - BM) * LDS + ch * 4, W + (long)(n0 + (row - BM)) * Kstride + k0 + ch * 4, pol_a);
                    }
                }
                cp_async_commit();
            }
        }
        // Drain remaining iterations (compile-time buffer indices).
        {
            #pragma unroll
            for (int st = 0; st < STAGES; ++st) {
                if (i + st >= iters) break;
                cp_async_wait<STAGES - 2>();
                __syncthreads();
                const float* smA = smem + st * STAGE_FLOATS;
                const float* smW = smA + BM * LDS;
                const float* a_base = smA + warpM * 64 * LDS + g * LDS + tg * 2;
                const float* b_base = smW + warpN * 32 * LDS + tg * 2 + g * LDS;
                #pragma unroll
                for (int kk = 0; kk < BK / 8; ++kk) {
                    float fa[4][4];
                    float fb[4][2];
                    #pragma unroll
                    for (int mi = 0; mi < 4; ++mi) {
                        const float* p = a_base + kk * 8 + mi * 16 * LDS;
                        float2 v01 = *reinterpret_cast<const float2*>(p);
                        float2 v23 = *reinterpret_cast<const float2*>(p + 8 * LDS);
                        fa[mi][0] = v01.x; fa[mi][1] = v23.x;
                        fa[mi][2] = v01.y; fa[mi][3] = v23.y;
                    }
                    #pragma unroll
                    for (int ni = 0; ni < 4; ++ni) {
                        float2 v01 = *reinterpret_cast<const float2*>(b_base + kk * 8 + ni * 16 * LDS);
                        fb[ni][0] = v01.x; fb[ni][1] = v01.y;
                    }
                    #pragma unroll
                    for (int mi = 0; mi < 4; ++mi)
                        #pragma unroll
                        for (int ni = 0; ni < 4; ++ni)
                            mma_tf32(acc[mi][ni][0], acc[mi][ni][1], acc[mi][ni][2], acc[mi][ni][3],
                                     fa[mi][0], fa[mi][1], fa[mi][2], fa[mi][3],
                                     fb[ni][0], fb[ni][1]);
                }
                cp_async_commit();
            }
        }

        cp_async_wait<0>();
        __syncthreads();

        // Stage the 128x128 partial tile through smem for coalesced float4
        // global writes.
        float* smC = smem;
        #pragma unroll
        for (int mi = 0; mi < 4; ++mi) {
            const int r0 = warpM * 64 + mi * 16;
            #pragma unroll
            for (int ni = 0; ni < 4; ++ni) {
                const int c0 = warpN * 32 + ni * 8;
                float* dst0 = smC + (r0 + g) * (BN + 4) + c0 + tg * 2;
                float* dst1 = smC + (r0 + g + 8) * (BN + 4) + c0 + tg * 2;
                *reinterpret_cast<float2*>(dst0) = make_float2(acc[mi][ni][0], acc[mi][ni][1]);
                *reinterpret_cast<float2*>(dst1) = make_float2(acc[mi][ni][2], acc[mi][ni][3]);
            }
        }
        __syncthreads();

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

// Split-K reduction with fused epilogue:
// out = min(sum_s ws[s] + bias, c) - c   (vectorized float4, smem-cached bias)
__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) {
    extern __shared__ float4 sh_bias[];
    for (long i = threadIdx.x; i < biasN4; i += blockDim.x) sh_bias[i] = bias[i];
    __syncthreads();

    const long stride = (long)gridDim.x * blockDim.x;
    for (long i = blockIdx.x * (long)blockDim.x + threadIdx.x; i < n4; i += stride) {
        float4 acc = sh_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 32");
    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;

    // Cost model over split-K factor S.
    const double hbm_bw   = 1.8e12;
    const double l2_bw    = 5.0e12;
    const double flop_sm  = 1.0e12;   // sustained TF32 per SM (fraction of peak)
    const double w_bytes  = (double)N * (double)K * 4.0;
    const double a_bytes  = (double)M * (double)K * 4.0;
    const double tile_mb  = (double)M * (double)N * 4.0;

    int bestS = 1;
    double bestT = 1e300;
    const int maxS = std::min(16, total_iters);
    for (int s = 1; s <= maxS; ++s) {
        const long items = (long)gridN * s;
        const double kiters = (double)total_iters / s;
        const double flop_item = 2.0 * BM * BN * (kiters * BK);
        const double t_comp = flop_item / flop_sm;
        const double t_w = (kiters * BK * BN * 4.0) / (hbm_bw / sms);
        const double t_a = (kiters * BK * BM * 4.0) / (l2_bw / sms);
        double t_item = std::max(t_comp, std::max(t_w, t_a)) + 20e-6;
        const double t = ((double)(items + sms - 1) / sms) * t_item
                         + (s > 1 ? (double)s * tile_mb * 2.0 / hbm_bw : 0.0);
        if (t < bestT - 1e-12) { bestT = t; bestS = s; }
    }
    const int S = bestS;
    const int kiters_per = total_iters / S;   // caller guarantees divisibility
    const long total_items = (long)gridN * S;

    auto out = at::empty({M, N}, x.options());
    auto ws = at::empty({(long)S, M, N}, x.options());
    float* ws_ptr = ws.data_ptr<float>();

    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();
    const int ggrid = (int)std::min<long>((long)sms, total_items);
    gemm_tf32_splitk_kernel<<<ggrid, NTHREADS, SMEM_BYTES, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), ws_ptr,
        gridN, (int)K, (int)N, kiters_per, (int)total_items);

    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, 2048L);
    reduce_min_sub_kernel<<<rgrid, NTHREADS, biasN4 * 16, stream>>>(
        reinterpret_cast<const float4*>(ws_ptr),
        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_min_sub_v3",
    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.

    For the benchmark shape everything (GEMM + bias + min + subtract) runs on
    a custom persistent split-K TF32 tensor-core kernel:
      * the 1 GiB FP32 weight matrix is streamed from HBM exactly once via a
        5-stage cp.async pipeline,
      * the 8 MiB activation matrix is pinned in L2 (evict_last) and served
        to all blocks from cache,
      * mma.sync.m16n8k8 with register-resident accumulators (no wmma shared
        round-trips), register double-buffered fragments,
      * the epilogue (bias add, torch.min, subtract) is fused into the split-K
        reduction so the final output is written exactly once.

    Falls back to the standard PyTorch path (TF32 cuBLAS) for any other
    shape/device.
    """
    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.dim() == 2
                and x.size(0) == 128
                and self.linear.in_features % 32 == 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
```