```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 for C = tril(A @ B) where A and B are lower triangular.
#
# Exploits triangular structure:
#   * The product of two lower-triangular matrices is lower triangular, so
#     blocks above the diagonal are just zeroed (no FLOPs).
#   * For a lower-triangular output block, only k-tiles in
#     [col_block, row_block] contribute (A[i,k]=0 for k>i, B[k,j]=0 for k<j).
# This reduces the total work from O(N^3) to ~N^3/6 (~6x fewer FLOPs), and
# fuses the tril() mask into the same kernel (no second pass over memory).
tril_matmul_cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cstdint>

#define BM 128
#define BN 128
#define BK 16
#define TM 8
#define TN 8
#define LDS (BM + 4)   // padded shared-memory stride (132)

__global__ void tril_matmul_kernel(const float* __restrict__ A,
                                   const float* __restrict__ B,
                                   float* __restrict__ C,
                                   const int N) {
    __shared__ float As[BK][LDS];
    __shared__ float Bs[BK][LDS];

    const int br  = blockIdx.y;          // block row of C
    const int bc  = blockIdx.x;          // block col of C
    const int tid = threadIdx.x;         // 0..255
    const int tx  = tid % (BN / TN);     // 0..15
    const int ty  = tid / (BN / TN);     // 0..15

    const int row0 = br * BM + ty * TM;  // this thread's first output row
    const int col0 = bc * BN + tx * TN;  // this thread's first output col

    // Strictly upper-triangular block: result is exactly zero, no math needed.
    if (bc > br) {
        #pragma unroll
        for (int i = 0; i < TM; ++i) {
            const int r = row0 + i;
            if (r < N) {
                #pragma unroll
                for (int j = 0; j < TN; ++j) {
                    const int c = col0 + j;
                    if (c < N) C[(size_t)r * N + c] = 0.0f;
                }
            }
        }
        return;
    }

    float acc[TM][TN];
    #pragma unroll
    for (int i = 0; i < TM; ++i)
        #pragma unroll
        for (int j = 0; j < TN; ++j)
            acc[i][j] = 0.0f;

    // k-tile range: B[k,j] != 0 only for k >= bc*BN, A[i,k] != 0 only for
    // k < (br+1)*BM. Everything outside contributes exactly zero.
    const int t0   = bc * (BN / BK);
    int       t1   = ((br + 1) * BM - 1) / BK;
    const int tmax = (N - 1) / BK;
    if (t1 > tmax) t1 = tmax;

    const float4* A4 = reinterpret_cast<const float4*>(A);
    const float4* B4 = reinterpret_cast<const float4*>(B);

    for (int t = t0; t <= t1; ++t) {
        // Load A tile (BM x BK), transposed into As. 512 float4 loads / 256 threads.
        #pragma unroll
        for (int i = 0; i < 2; ++i) {
            const int p  = tid + i * 256;   // 0..511
            const int m  = p >> 2;          // tile row 0..127
            const int kg = p & 3;           // float4 group along k
            const int gr = br * BM + m;
            const int gc = t * BK + kg * 4;
            float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
            if (gr < N && gc < N) v = A4[((size_t)gr * N + gc) >> 2];
            As[kg * 4 + 0][m] = v.x;
            As[kg * 4 + 1][m] = v.y;
            As[kg * 4 + 2][m] = v.z;
            As[kg * 4 + 3][m] = v.w;
        }
        // Load B tile (BK x BN) into Bs.
        #pragma unroll
        for (int i = 0; i < 2; ++i) {
            const int p    = tid + i * 256; // 0..511
            const int krow = p >> 5;        // 0..15
            const int cg   = p & 31;        // float4 group along n
            const int gr   = t * BK + krow;
            const int gc   = bc * BN + cg * 4;
            float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
            if (gr < N && gc < N) v = B4[((size_t)gr * N + gc) >> 2];
            *reinterpret_cast<float4*>(&Bs[krow][cg * 4]) = v;
        }

        __syncthreads();

        #pragma unroll
        for (int k = 0; k < BK; ++k) {
            float rm[TM], rn[TN];
            *reinterpret_cast<float4*>(&rm[0]) =
                *reinterpret_cast<const float4*>(&As[k][ty * TM]);
            *reinterpret_cast<float4*>(&rm[4]) =
                *reinterpret_cast<const float4*>(&As[k][ty * TM + 4]);
            *reinterpret_cast<float4*>(&rn[0]) =
                *reinterpret_cast<const float4*>(&Bs[k][tx * TN]);
            *reinterpret_cast<float4*>(&rn[4]) =
                *reinterpret_cast<const float4*>(&Bs[k][tx * TN + 4]);
            #pragma unroll
            for (int i = 0; i < TM; ++i)
                #pragma unroll
                for (int j = 0; j < TN; ++j)
                    acc[i][j] += rm[i] * rn[j];
        }
        __syncthreads();
    }

    // Write back with fused tril mask.
    #pragma unroll
    for (int i = 0; i < TM; ++i) {
        const int r = row0 + i;
        if (r < N) {
            #pragma unroll
            for (int j = 0; j < TN; ++j) {
                const int c = col0 + j;
                if (c < N) C[(size_t)r * N + c] = (c <= r) ? acc[i][j] : 0.0f;
            }
        }
    }
}

torch::Tensor tril_matmul_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(A.scalar_type() == torch::kFloat32, "Inputs must be float32");
    const int N = (int)A.size(0);

    auto C = torch::empty_like(A);

    dim3 block(256);
    dim3 grid((N + BN - 1) / BN, (N + BM - 1) / BM);

    tril_matmul_kernel<<<grid, block>>>(
        A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), N);

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
"""

tril_matmul_cpp_source = (
    "torch::Tensor tril_matmul_cuda(torch::Tensor A, torch::Tensor B);"
)

tril_matmul_ext = load_inline(
    name="tril_matmul_ext",
    cpp_sources=tril_matmul_cpp_source,
    cuda_sources=tril_matmul_cuda_source,
    functions=["tril_matmul_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version: fused lower-triangular matmul (C = tril(A @ B)) as a
    single custom CUDA kernel that skips all provably-zero work (~6x fewer
    FLOPs than dense matmul) and applies the tril mask in-kernel.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.op = tril_matmul_ext

    def forward(self, A, B):
        if (
            A.is_cuda
            and A.dtype == torch.float32
            and A.dim() == 2
            and A.size(0) == A.size(1)
            and B.shape == A.shape
            and A.size(0) % 4 == 0  # float4 alignment requirement
        ):
            return self.op.tril_matmul_cuda(A.contiguous(), B.contiguous())
        # Safe fallback for anything unexpected.
        return torch.tril(torch.matmul(A, B))
```