```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

# Fused lower-triangular matmul (C = tril(A @ B)) using TF32 tensor cores.
#
# Key ideas:
#  * A and B are lower triangular, so C is lower triangular and the tril() is
#    redundant. In 128x128 tile space, C tile (t, s) only receives
#    contributions from k tiles s..t; all other block products are exactly
#    zero and are skipped (~6x FLOP reduction vs. dense GEMM).
#  * The remaining tiles are computed with WMMA TF32 tensor-core MMAs
#    (matching the numerics of the TF32-enabled cuBLAS reference), with a
#    cp.async double-buffered shared-memory pipeline.
#  * Strictly-upper C tiles are written as zeros directly by the kernel, so
#    no separate memset/fill pass is needed.
_tril_mm_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cuda_pipeline.h>
#include <mma.h>

using namespace nvcuda;

#define BM 128
#define BN 128
#define BK 32
#define LDSA (BK + 8)          // padded strides (32B aligned, bank-conflict friendly)
#define LDSB (BN + 8)
#define NTHREADS 256
#define WARPS_M 2
#define WARPS_N 4
#define WM (BM / WARPS_M)      // 64
#define WN (BN / WARPS_N)      // 32
#define SMEM_FLOATS (2 * (BM * LDSA + BK * LDSB))

__global__ void __launch_bounds__(NTHREADS, 2)
tril_matmul_tf32_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)
    const int tid = threadIdx.x;
    float* Cg = C + (size_t)t * BM * N + (size_t)s * BN;

    if (s > t) {
        // Strictly upper-triangular output tile: exactly zero.
        #pragma unroll
        for (int i = 0; i < (BM * BN / 4) / NTHREADS; ++i) {
            const int idx = tid + i * NTHREADS;
            const int r = idx >> 5;
            const int c4 = (idx & 31) * 4;
            *reinterpret_cast<float4*>(Cg + (size_t)r * N + c4) =
                make_float4(0.f, 0.f, 0.f, 0.f);
        }
        return;
    }

    extern __shared__ float smem[];
    float* As[2] = {smem, smem + BM * LDSA};
    float* Bs[2] = {smem + 2 * BM * LDSA, smem + 2 * BM * LDSA + BK * LDSB};

    const int warpId = tid >> 5;
    const int warpM = warpId / WARPS_N;
    const int warpN = warpId % WARPS_N;

    const float* Ag = A + (size_t)t * BM * N;   // A tile rows
    const float* Bg = B + (size_t)s * BN;       // B tile cols

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

    // Global -> shared loader coordinates (16B cp.async chunks).
    const int aRow = tid >> 3;        // rows aRow + i*32 cover 0..127
    const int aC4  = (tid & 7) * 4;   // 8 chunks per 32-col row
    const int bRow = tid >> 5;        // rows bRow + i*8 cover 0..31
    const int bC4  = (tid & 31) * 4;  // 32 chunks per 128-col row

    auto issue_tile = [&](int kt, int buf) {
        const int k0 = kStart + kt * BK;
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            const int row = aRow + i * 32;
            __pipeline_memcpy_async(&As[buf][row * LDSA + aC4],
                                    Ag + (size_t)row * N + k0 + aC4, 16);
        }
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            const int row = bRow + i * 8;
            __pipeline_memcpy_async(&Bs[buf][row * LDSB + bC4],
                                    Bg + (size_t)(k0 + row) * N + bC4, 16);
        }
    };

    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc[WM / 16][WN / 16];
    #pragma unroll
    for (int i = 0; i < WM / 16; ++i)
        #pragma unroll
        for (int j = 0; j < WN / 16; ++j)
            wmma::fill_fragment(acc[i][j], 0.0f);

    issue_tile(0, 0);
    __pipeline_commit();

    for (int kt = 0; kt < numKt; ++kt) {
        if (kt + 1 < numKt) issue_tile(kt + 1, (kt + 1) & 1);
        __pipeline_commit();
        __pipeline_wait_prior(1);
        __syncthreads();

        const float* Asb = As[kt & 1];
        const float* Bsb = Bs[kt & 1];

        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            wmma::fragment<wmma::matrix_a, 16, 16, 8,
                           wmma::precision::tf32, wmma::row_major> af[WM / 16];
            wmma::fragment<wmma::matrix_b, 16, 16, 8,
                           wmma::precision::tf32, wmma::row_major> bf[WN / 16];
            #pragma unroll
            for (int i = 0; i < WM / 16; ++i)
                wmma::load_matrix_sync(af[i], &Asb[(warpM * WM + i * 16) * LDSA + kk], LDSA);
            #pragma unroll
            for (int j = 0; j < WN / 16; ++j)
                wmma::load_matrix_sync(bf[j], &Bsb[kk * LDSB + warpN * WN + j * 16], LDSB);
            #pragma unroll
            for (int i = 0; i < WM / 16; ++i)
                #pragma unroll
                for (int j = 0; j < WN / 16; ++j)
                    wmma::mma_sync(acc[i][j], af[i], bf[j], acc[i][j]);
        }
        __syncthreads();
    }

    // Epilogue: write the full 128x128 tile. Upper entries of diagonal tiles
    // are exact zeros (all their products are exactly 0), matching tril.
    #pragma unroll
    for (int i = 0; i < WM / 16; ++i)
        #pragma unroll
        for (int j = 0; j < WN / 16; ++j)
            wmma::store_matrix_sync(
                &Cg[(size_t)(warpM * WM + i * 16) * N + warpN * WN + j * 16],
                acc[i][j], N, wmma::mem_row_major);
}

torch::Tensor tril_matmul(torch::Tensor A, torch::Tensor B) {
    const int N = (int)A.size(0);
    auto C = torch::empty({N, N}, A.options());
    const int NT = N / BM;
    dim3 grid(NT, NT);
    dim3 block(NTHREADS);
    const int smemBytes = (int)(SMEM_FLOATS * sizeof(float));
    cudaFuncSetAttribute(tril_matmul_tf32_kernel,
                         cudaFuncAttributeMaxDynamicSharedMemorySize, smemBytes);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    tril_matmul_tf32_kernel<<<grid, block, smemBytes, 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_tf32_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
    TF32 tensor-core 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()
            and torch.cuda.get_device_capability(A.device)[0] >= 8
        ):
            return _tril_mm.tril_matmul(A, B)
        return torch.tril(torch.matmul(A, B))
```