```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.
#
# Optimizations vs. the previous version:
#  * Block tile 128x128, but now 4 warps with a 64x64 warp tile. This doubles
#    the MMA:instruction ratio (16 MMAs per 8 fragment loads per k-step), so
#    the tensor cores become the bottleneck instead of the instruction issue
#    path (the old 64x32 warp tile was issue/latency bound at ~33% of peak).
#  * Bank-conflict-free shared-memory layouts (LDSA=36, LDSB=132 words give a
#    4-word row-stride modulo 32 banks -> conflict-free ldmatrix for both A
#    and B fragments while keeping 16B alignment for cp.async).
#  * cp.async double-buffered global->shared pipeline (2 stages, 70.7KB),
#    2 thread blocks per SM on A100 (141KB smem, ~190 regs/thread).
#  * Block-triangular work skipping: C tile (t, s) only accumulates k-tiles
#    s..t (~6x FLOP reduction vs. dense GEMM); strictly-upper tiles are
#    written as zeros directly; heavy rows are scheduled first so the tail is
#    made of the lightest blocks.
_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 36            // padded strides: conflict-free ldmatrix, 16B aligned
#define LDSB 132
#define NTHREADS 128       // 4 warps, each with a 64x64 warp tile
#define STAGE_A (BM * LDSA)
#define STAGE_B (BK * LDSB)
#define SMEM_FLOATS (2 * (STAGE_A + STAGE_B))

__global__ void __launch_bounds__(NTHREADS, 2)
tril_matmul_tf32_v2_kernel(const float* __restrict__ A,
                           const float* __restrict__ B,
                           float* __restrict__ C,
                           const int N)
{
#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800)
    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 + STAGE_A};
    float* Bs[2] = {smem + 2 * STAGE_A, smem + 2 * STAGE_A + STAGE_B};

    const int warpId = tid >> 5;
    const int warpM = warpId >> 1;   // 0..1
    const int warpN = warpId & 1;    // 0..1

    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*16 cover 0..127
    const int aC4  = (tid & 7) * 4;   // 8 chunks per 32-col row
    const int bRow = tid >> 5;        // rows bRow + i*4 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 < 8; ++i) {
            const int row = aRow + i * 16;
            __pipeline_memcpy_async(&As[buf][row * LDSA + aC4],
                                    Ag + (size_t)row * N + k0 + aC4, 16);
        }
        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            const int row = bRow + i * 4;
            __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[4][4];
    #pragma unroll
    for (int i = 0; i < 4; ++i)
        #pragma unroll
        for (int j = 0; j < 4; ++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[4];
            wmma::fragment<wmma::matrix_b, 16, 16, 8,
                           wmma::precision::tf32, wmma::row_major> bf[4];
            #pragma unroll
            for (int i = 0; i < 4; ++i)
                wmma::load_matrix_sync(af[i], &Asb[(warpM * 64 + i * 16) * LDSA + kk], LDSA);
            #pragma unroll
            for (int j = 0; j < 4; ++j)
                wmma::load_matrix_sync(bf[j], &Bsb[kk * LDSB + warpN * 64 + j * 16], LDSB);
            #pragma unroll
            for (int i = 0; i < 4; ++i)
                #pragma unroll
                for (int j = 0; j < 4; ++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 < 4; ++i)
        #pragma unroll
        for (int j = 0; j < 4; ++j)
            wmma::store_matrix_sync(
                &Cg[(size_t)(warpM * 64 + i * 16) * N + warpN * 64 + j * 16],
                acc[i][j], N, wmma::mem_row_major);
#else
    // Portable fallback (never used; the host wrapper guards capability >= 8).
    const int s = blockIdx.x;
    const int t = gridDim.y - 1 - blockIdx.y;
    const int tid = threadIdx.x;
    float* Cg = C + (size_t)t * BM * N + (size_t)s * BN;
    for (int idx = tid; idx < BM * BN; idx += NTHREADS) {
        const int r = idx / BN, c = idx % BN;
        const int gr = t * BM + r, gc = s * BN + c;
        float v = 0.f;
        if (s <= t && gc <= gr) {
            for (int k = gc; k <= gr; ++k) v += A[(size_t)gr * N + k] * B[(size_t)k * N + gc];
        }
        Cg[(size_t)r * N + c] = v;
    }
#endif
}

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);
    static const int smemBytes = (int)(SMEM_FLOATS * sizeof(float));
    static bool configured = []() {
        cudaFuncSetAttribute(tril_matmul_tf32_v2_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, smemBytes);
        cudaFuncSetAttribute(tril_matmul_tf32_v2_kernel,
                             cudaFuncAttributePreferredSharedMemoryCarveout, 100);
        return true;
    }();
    (void)configured;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    tril_matmul_tf32_v2_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_v2_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))
```