```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 with cvt.rna.tf32.f32
#     (no slow wmma smem path), accumulator fragments kept resident in
#     registers across the whole K-chunk.
#   * Multi-stage cp.async pipeline (stage count chosen at runtime from the
#     device's max opt-in shared memory, so it runs on sm_80/86/90) with an
#     L2 evict_last hint for the small activation matrix.
#   * Persistent grid + split-K: (n-tile, k-chunk) work items are distributed
#     over a grid sized to the SM count, so there is no wave quantization tail.
#   * 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 LDS 40
#define NTHREADS 256
#define STAGE_FLOATS ((BM + BN) * LDS)
#define CPLDS (BK / 4)

__device__ __forceinline__ void cp_async16_pol(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_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));
}

__device__ __forceinline__ float to_tf32(float x) {
    unsigned u;
    asm volatile("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(u) : "f"(x));
    return __uint_as_float(u);
}

__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 strided range of work items; each item computes a
// 128x128 output tile over a K/S chunk and writes raw partials to ws[s].
template <int STAGES>
__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) * CPLDS;   // 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;
        const int iters = 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 (commit every stage so the
        // group accounting stays uniform even past the chunk end).
        #pragma unroll
        for (int st = 0; st < STAGES - 1; ++st) {
            if (st < iters) {
                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_pol(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);
                }
            }
            cp_async_commit();
        }

        // Mainloop in blocks of STAGES so the buffer index is compile-time.
        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 + g * LDS + tg * 2;

                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] = to_tf32(v01.x); fa[0][mi][1] = to_tf32(v23.x);
                    fa[0][mi][2] = to_tf32(v01.y); fa[0][mi][3] = to_tf32(v23.y);
                }
                #pragma unroll
                for (int ni = 0; ni < 4; ++ni) {
                    float2 v01 = *reinterpret_cast<const float2*>(b_base + ni * 8 * LDS);
                    fb[0][ni][0] = to_tf32(v01.x); fb[0][ni][1] = to_tf32(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] = to_tf32(v01.x); fa[(kk + 1) & 1][mi][1] = to_tf32(v23.x);
                            fa[(kk + 1) & 1][mi][2] = to_tf32(v01.y); fa[(kk + 1) & 1][mi][3] = to_tf32(v23.y);
                        }
                        #pragma unroll
                        for (int ni = 0; ni < 4; ++ni) {
                            float2 v01 = *reinterpret_cast<const float2*>(b_base + (kk + 1) * 8 + ni * 8 * LDS);
                            fb[(kk + 1) & 1][ni][0] = to_tf32(v01.x); fb[(kk + 1) & 1][ni][1] = to_tf32(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_pol(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);
                    }
                }
                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 + g * LDS + tg * 2;
                #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] = to_tf32(v01.x); fa[mi][1] = to_tf32(v23.x);
                        fa[mi][2] = to_tf32(v01.y); fa[mi][3] = to_tf32(v23.y);
                    }
                    #pragma unroll
                    for (int ni = 0; ni < 4; ++ni) {
                        float2 v01 = *reinterpret_cast<const float2*>(b_base + kk * 8 + ni * 8 * LDS);
                        fb[ni][0] = to_tf32(v01.x); fb[ni][1] = to_tf32(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; bias read
// through L2/L1 cache, no dynamic smem so no opt-in attribute is needed)
__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;
    }
}

template <int STAGES>
static void launch_gemm(const float* A, const float* W, float* ws,
                        int grid, int gridN, int K, int N,
                        int kiters_per, int total_items, cudaStream_t stream) {
    static bool attr_set = false;
    if (!attr_set) {
        cudaFuncSetAttribute(gemm_tf32_splitk_kernel<STAGES>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             STAGES * STAGE_FLOATS * 4);
        attr_set = true;
    }
    gemm_tf32_splitk_kernel<STAGES><<<grid, NTHREADS, STAGES * STAGE_FLOATS * 4, stream>>>(
        A, W, ws, gridN, K, N, kiters_per, total_items);
}

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;
    int max_smem = 0;
    cudaDeviceGetAttribute(&max_smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);

    // Cost model over split-K factor S (divisors of total_iters only).
    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 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) {
        if (total_iters % s != 0) continue;
        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;
    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>();

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int ggrid = (int)std::min<long>((long)sms, total_items);

    int stages = max_smem / (STAGE_FLOATS * 4);
    if (stages > 5) stages = 5;
    TORCH_CHECK(stages >= 2, "insufficient shared memory for pipelined kernel");
    switch (stages) {
        case 5: launch_gemm<5>(xc.data_ptr<float>(), wc.data_ptr<float>(), ws_ptr,
                               ggrid, gridN, (int)K, (int)N, kiters_per, (int)total_items, stream); break;
        case 4: launch_gemm<4>(xc.data_ptr<float>(), wc.data_ptr<float>(), ws_ptr,
                               ggrid, gridN, (int)K, (int)N, kiters_per, (int)total_items, stream); break;
        case 3: launch_gemm<3>(xc.data_ptr<float>(), wc.data_ptr<float>(), ws_ptr,
                               ggrid, gridN, (int)K, (int)N, kiters_per, (int)total_items, stream); break;
        default: launch_gemm<2>(xc.data_ptr<float>(), wc.data_ptr<float>(), ws_ptr,
                                ggrid, gridN, (int)K, (int)N, kiters_per, (int)total_items, stream); break;
    }

    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, 4096L);
    reduce_min_sub_kernel<<<rgrid, NTHREADS, 0, 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_v4",
    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
        multi-stage cp.async pipeline (stage count adapts to the device's
        shared-memory size),
      * 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 and 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.weight.is_cuda
                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
```