```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 CUDA kernel: block-triangular GEMM.
#
# A and B are lower-triangular, so C = A @ B is also lower-triangular and the
# torch.tril is mathematically redundant. In 128x128 tile space, tile (t, s) of
# C only receives contributions from A tiles (t, k) and B tiles (k, s) with
# s <= k <= t. Every other block product is exactly zero, so we simply skip it.
# This reduces the number of 128^3 tile-GEMMs from NT^3 to C(NT+2, 3) ~= NT^3/6
# (~6x FLOP reduction), while a register-blocked, double-buffered SGEMM tile
# pipeline keeps per-FLOP efficiency close to cuBLAS. The tril is fused for
# free: only tiles with t >= s are computed and the output is pre-zeroed.
_tril_mm_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

#define BM 128
#define BN 128
#define BK 16
#define TM 8
#define TN 8
#define NTHREADS ((BM / TM) * (BN / TN))  // 256

__global__ void __launch_bounds__(NTHREADS)
tril_matmul_kernel(const float* __restrict__ A,
                   const float* __restrict__ B,
                   float* __restrict__ C,
                   const int N)
{
    const int s = blockIdx.x;                  // column tile of C
    const int t = gridDim.y - 1 - blockIdx.y;  // row tile of C (heavy rows first)
    if (s > t) return;                         // strictly-upper tiles stay zero

    __shared__ float As[BK * (BM + 4)];        // transposed A staging, padded
    __shared__ float Bs[BK * BN];

    const int tid = threadIdx.x;
    const int threadRow = tid / (BN / TN);     // 0..15
    const int threadCol = tid % (BN / TN);     // 0..15

    const float* __restrict__ Ag = A + (size_t)t * BM * N;
    const float* __restrict__ Bg = B + (size_t)s * BN;
    float* __restrict__ Cg = C + (size_t)t * BM * N + (size_t)s * BN;

    // Global->shared loader coordinates (float4 vectorized).
    const int innerRowA = tid / (BK / 4);      // 0..63  (2 passes cover 128 rows)
    const int innerColA = tid % (BK / 4);      // 0..3
    const int innerRowB = tid / (BN / 4);      // 0..7   (2 passes cover 16 rows)
    const int innerColB = tid % (BN / 4);      // 0..31

    float acc[TM][TN] = {{0.0f}};

    const int kStart = s * BN;                 // first contributing k element
    const int numKt = (t - s + 1) * (BN / BK); // k-tiles: s..t in 128-tile space

    float4 regA[BM / 64];
    float4 regB[BK / 8];

    // Prologue: load first k-tile into shared memory.
    {
        const int k0 = kStart;
        #pragma unroll
        for (int off = 0; off < BM; off += 64) {
            const float4 tmp = *reinterpret_cast<const float4*>(
                Ag + (size_t)(innerRowA + off) * N + k0 + innerColA * 4);
            As[(innerColA * 4 + 0) * (BM + 4) + innerRowA + off] = tmp.x;
            As[(innerColA * 4 + 1) * (BM + 4) + innerRowA + off] = tmp.y;
            As[(innerColA * 4 + 2) * (BM + 4) + innerRowA + off] = tmp.z;
            As[(innerColA * 4 + 3) * (BM + 4) + innerRowA + off] = tmp.w;
        }
        #pragma unroll
        for (int off = 0; off < BK; off += 8) {
            const float4 tmp = *reinterpret_cast<const float4*>(
                Bg + (size_t)(innerRowB + off) * N + k0 + innerColB * 4);
            *reinterpret_cast<float4*>(&Bs[(innerRowB + off) * BN + innerColB * 4]) = tmp;
        }
    }
    __syncthreads();

    for (int kt = 0; kt < numKt; ++kt) {
        // Prefetch next k-tile from global memory into registers.
        if (kt + 1 < numKt) {
            const int k0 = kStart + (kt + 1) * BK;
            #pragma unroll
            for (int off = 0; off < BM; off += 64) {
                regA[off / 64] = *reinterpret_cast<const float4*>(
                    Ag + (size_t)(innerRowA + off) * N + k0 + innerColA * 4);
            }
            #pragma unroll
            for (int off = 0; off < BK; off += 8) {
                regB[off / 8] = *reinterpret_cast<const float4*>(
                    Bg + (size_t)(innerRowB + off) * N + k0 + innerColB * 4);
            }
        }

        // Compute 8x8 outer-product micro-tiles from shared memory.
        #pragma unroll
        for (int k = 0; k < BK; ++k) {
            float afrag[TM];
            float bfrag[TN];
            *reinterpret_cast<float4*>(&afrag[0]) =
                *reinterpret_cast<const float4*>(&As[k * (BM + 4) + threadRow * TM]);
            *reinterpret_cast<float4*>(&afrag[4]) =
                *reinterpret_cast<const float4*>(&As[k * (BM + 4) + threadRow * TM + 4]);
            *reinterpret_cast<float4*>(&bfrag[0]) =
                *reinterpret_cast<const float4*>(&Bs[k * BN + threadCol * TN]);
            *reinterpret_cast<float4*>(&bfrag[4]) =
                *reinterpret_cast<const float4*>(&Bs[k * BN + threadCol * TN + 4]);
            #pragma unroll
            for (int i = 0; i < TM; ++i) {
                #pragma unroll
                for (int j = 0; j < TN; ++j) {
                    acc[i][j] = fmaf(afrag[i], bfrag[j], acc[i][j]);
                }
            }
        }
        __syncthreads();

        // Commit prefetched tile to shared memory.
        if (kt + 1 < numKt) {
            #pragma unroll
            for (int off = 0; off < BM; off += 64) {
                const float4 tmp = regA[off / 64];
                As[(innerColA * 4 + 0) * (BM + 4) + innerRowA + off] = tmp.x;
                As[(innerColA * 4 + 1) * (BM + 4) + innerRowA + off] = tmp.y;
                As[(innerColA * 4 + 2) * (BM + 4) + innerRowA + off] = tmp.z;
                As[(innerColA * 4 + 3) * (BM + 4) + innerRowA + off] = tmp.w;
            }
            #pragma unroll
            for (int off = 0; off < BK; off += 8) {
                *reinterpret_cast<float4*>(&Bs[(innerRowB + off) * BN + innerColB * 4]) = regB[off / 8];
            }
        }
        __syncthreads();
    }

    // Epilogue: write the full 128x128 tile (upper entries of diagonal tiles
    // are exact zeros, matching torch.tril semantics).
    #pragma unroll
    for (int i = 0; i < TM; ++i) {
        const int row = threadRow * TM + i;
        #pragma unroll
        for (int j = 0; j < TN; j += 4) {
            *reinterpret_cast<float4*>(Cg + (size_t)row * N + threadCol * TN + j) =
                make_float4(acc[i][j], acc[i][j + 1], acc[i][j + 2], acc[i][j + 3]);
        }
    }
}

torch::Tensor tril_matmul(torch::Tensor A, torch::Tensor B) {
    const int N = (int)A.size(0);
    auto C = torch::zeros({N, N}, A.options());
    const int NT = N / BM;
    dim3 grid(NT, NT);
    dim3 block(NTHREADS);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    tril_matmul_kernel<<<grid, block, 0, stream>>>(
        A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), N);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
"""

_tril_mm_cpp_src = "torch::Tensor tril_matmul(torch::Tensor A, torch::Tensor B);"

_tril_mm = load_inline(
    name="tril_matmul_ext",
    cpp_sources=_tril_mm_cpp_src,
    cuda_sources=_tril_mm_cuda_src,
    functions=["tril_matmul"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized model: fused lower-triangular matmul via a custom block-triangular
    CUDA GEMM kernel. Falls back to the reference PyTorch ops for any input the
    kernel does not support.
    """
    def __init__(self):
        super(ModelNew, self).__init__()

    def forward(self, A, B):
        if (
            A.is_cuda and B.is_cuda
            and A.dtype == torch.float32 and B.dtype == torch.float32
            and A.dim() == 2 and B.dim() == 2
            and A.shape == B.shape
            and A.size(0) == A.size(1)
            and A.size(0) > 0
            and A.size(0) % 128 == 0
            and A.is_contiguous() and B.is_contiguous()
        ):
            return _tril_mm.tril_matmul(A, B)
        return torch.tril(torch.matmul(A, B))
```