```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) with lower-triangular A and B.
#
# Key optimizations vs. torch.matmul + torch.tril:
#   * Only the lower-triangular 128x128 tiles of C are computed (the product of
#     two lower-triangular matrices is lower-triangular), skipping ~1/2 of tiles.
#   * For each output tile, the k-reduction only spans the range that can
#     contribute non-zero terms, reducing total FLOPs from O(N^3) to ~N^3/6.
#   * Register-blocked (8x8 per thread) tiled GEMM with shared-memory staging,
#     software pipelining (next tile prefetched to registers while computing),
#     vectorized float4 global/shared accesses, and the final `tril` masking is
#     fully fused into the kernel.
cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

#define BM 128
#define BN 128
#define BK 16
#define TM 8
#define TN 8
#define NTHREADS 256

// Load a (BM x BK) slab of A and a (BK x BN) slab of B into registers.
// Elements above the diagonal are zeroed so the matrices are treated as
// exactly lower-triangular.
__device__ __forceinline__ void load_slab_regs(
    const float* __restrict__ A, const float* __restrict__ B,
    int N, int R0, int C0, int k0,
    int aRow, int aCol, int bRow, int bCol,
    float4 aReg[2], float4 bReg[2], bool vecOK)
{
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        int gr = R0 + aRow + (i << 6);   // global row of A slab
        int gc = k0 + aCol;              // global col of A slab
        float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
        if (gr < N) {
            if (vecOK && gc + 3 < N && gc + 3 <= gr) {
                v = *reinterpret_cast<const float4*>(A + (long)gr * N + gc);
            } else {
                const float* p = A + (long)gr * N + gc;
                v.x = (gc + 0 < N && gc + 0 <= gr) ? p[0] : 0.f;
                v.y = (gc + 1 < N && gc + 1 <= gr) ? p[1] : 0.f;
                v.z = (gc + 2 < N && gc + 2 <= gr) ? p[2] : 0.f;
                v.w = (gc + 3 < N && gc + 3 <= gr) ? p[3] : 0.f;
            }
        }
        aReg[i] = v;
    }
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        int gr = k0 + bRow + (i << 3);   // global row of B slab
        int gc = C0 + bCol;              // global col of B slab
        float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
        if (gr < N) {
            if (vecOK && gc + 3 < N && gr >= gc + 3) {
                v = *reinterpret_cast<const float4*>(B + (long)gr * N + gc);
            } else {
                const float* p = B + (long)gr * N + gc;
                v.x = (gc + 0 < N && gr >= gc + 0) ? p[0] : 0.f;
                v.y = (gc + 1 < N && gr >= gc + 1) ? p[1] : 0.f;
                v.z = (gc + 2 < N && gr >= gc + 2) ? p[2] : 0.f;
                v.w = (gc + 3 < N && gr >= gc + 3) ? p[3] : 0.f;
            }
        }
        bReg[i] = v;
    }
}

__device__ __forceinline__ void store_slab_shared(
    float As[BK][BM], float Bs[BK][BN],
    int aRow, int aCol, int bRow, int bCol,
    const float4 aReg[2], const float4 bReg[2])
{
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        int r = aRow + (i << 6);
        As[aCol + 0][r] = aReg[i].x;
        As[aCol + 1][r] = aReg[i].y;
        As[aCol + 2][r] = aReg[i].z;
        As[aCol + 3][r] = aReg[i].w;
    }
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        int r = bRow + (i << 3);
        *reinterpret_cast<float4*>(&Bs[r][bCol]) = bReg[i];
    }
}

__global__ void __launch_bounds__(NTHREADS, 2)
tril_matmul_kernel(const float* __restrict__ A,
                   const float* __restrict__ B,
                   float* __restrict__ C,
                   int N)
{
    __shared__ float As[BK][BM];
    __shared__ float Bs[BK][BN];

    const int numTiles = (N + BM - 1) / BM;
    const int totalBlocks = numTiles * (numTiles + 1) / 2;
    // Enumerate lower-triangular tiles in reverse order so the heaviest tiles
    // (far from the diagonal) are scheduled first.
    int bid = totalBlocks - 1 - (int)blockIdx.x;
    int rowTile = (int)((sqrtf(8.0f * (float)bid + 1.0f) - 1.0f) * 0.5f);
    while (((rowTile + 1) * (rowTile + 2)) / 2 <= bid) ++rowTile;
    while ((rowTile * (rowTile + 1)) / 2 > bid) --rowTile;
    int colTile = bid - (rowTile * (rowTile + 1)) / 2;

    const int R0 = rowTile * BM;
    const int C0 = colTile * BN;

    const int tid = threadIdx.x;
    const int threadRow = tid >> 4;   // 0..15
    const int threadCol = tid & 15;   // 0..15

    // Slab load coordinates (each thread moves 2 x float4 per matrix).
    const int aRow = tid >> 2;        // 0..63  (rows of A slab, stride 64)
    const int aCol = (tid & 3) << 2;  // 0,4,8,12
    const int bRow = tid >> 5;        // 0..7   (rows of B slab, stride 8)
    const int bCol = (tid & 31) << 2; // 0..124 step 4

    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;

    const bool vecOK = ((N & 3) == 0);
    const int kEnd = min(R0 + BM - 1, N - 1);

    float4 aReg[2], bReg[2];
    load_slab_regs(A, B, N, R0, C0, C0, aRow, aCol, bRow, bCol, aReg, bReg, vecOK);

    for (int k0 = C0; k0 <= kEnd; k0 += BK) {
        store_slab_shared(As, Bs, aRow, aCol, bRow, bCol, aReg, bReg);
        __syncthreads();

        const int kNext = k0 + BK;
        if (kNext <= kEnd) {
            load_slab_regs(A, B, N, R0, C0, kNext, aRow, aCol, bRow, bCol,
                           aReg, bReg, vecOK);
        }

        #pragma unroll
        for (int kk = 0; kk < BK; ++kk) {
            float ra[TM], rb[TN];
            #pragma unroll
            for (int i = 0; i < TM; ++i) ra[i] = As[kk][threadRow + (i << 4)];
            #pragma unroll
            for (int j = 0; j < TN; ++j) rb[j] = Bs[kk][threadCol + (j << 4)];
            #pragma unroll
            for (int i = 0; i < TM; ++i)
                #pragma unroll
                for (int j = 0; j < TN; ++j)
                    acc[i][j] = fmaf(ra[i], rb[j], acc[i][j]);
        }
        __syncthreads();
    }

    // Write the lower-triangular part of this C tile (upper part stays zero).
    #pragma unroll
    for (int i = 0; i < TM; ++i) {
        int gr = R0 + threadRow + (i << 4);
        if (gr >= N) continue;
        #pragma unroll
        for (int j = 0; j < TN; ++j) {
            int gc = C0 + threadCol + (j << 4);
            if (gc < N && gr >= gc) {
                C[(long)gr * N + gc] = acc[i][j];
            }
        }
    }
}

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 && B.scalar_type() == torch::kFloat32,
                "Only float32 tensors are supported");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Inputs must be 2-D");
    TORCH_CHECK(A.size(0) == A.size(1) && B.size(0) == B.size(1) && A.size(0) == B.size(0),
                "Inputs must be square matrices of the same size");

    auto Ac = A.contiguous();
    auto Bc = B.contiguous();
    const int N = (int)Ac.size(0);
    auto C = torch::zeros_like(Ac);
    if (N == 0) return C;

    const int numTiles = (N + BM - 1) / BM;
    const int totalBlocks = numTiles * (numTiles + 1) / 2;
    tril_matmul_kernel<<<totalBlocks, NTHREADS>>>(
        Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(), N);
    return C;
}
"""

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

tril_matmul = load_inline(
    name="tril_matmul",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["tril_matmul_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model that computes C = tril(A @ B) for lower-triangular A and B
    with a single fused custom CUDA kernel.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.tril_matmul = tril_matmul

    def forward(self, A, B):
        return self.tril_matmul.tril_matmul_cuda(A, B)
```