```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: fused lower-triangular matrix multiply (C = tril(A @ B)).
#
# Since A and B are lower triangular, C[i][j] = sum_{k=j}^{i} A[i][k]*B[k][j],
# which is only ~N^3/6 multiply-adds instead of N^3 for a dense GEMM (~5.5x less
# work). The kernel computes C in 128x128 tiles, launching only tiles in the
# lower triangle and restricting the k-loop to the non-zero range of each tile.
# Each block uses a register-blocked 8x8 micro-kernel with float4 vectorized
# global/shared memory accesses and conflict-free padded shared memory.
# The triangular mask is applied on write (fused tril), so no extra pass is needed.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define BM 128
#define BN 128
#define BK 8
#define TM 8
#define TN 8

__global__ void __launch_bounds__(256) lower_tri_mm_kernel(
        const float* __restrict__ A,
        const float* __restrict__ B,
        float* __restrict__ C,
        int N) {
    __shared__ __align__(16) float As[BK][BM + 4];
    __shared__ __align__(16) float Bs[BK][BN + 4];

    const int tid = threadIdx.x;

    // Map the linear block index to a lower-triangular tile (bi, bj), bi >= bj.
    // Launch in reverse order so the longest-running tiles start first.
    int t = (int)gridDim.x - 1 - (int)blockIdx.x;
    int bi = (int)((sqrtf(8.0f * (float)t + 1.0f) - 1.0f) * 0.5f);
    while (((bi + 1) * (bi + 2)) / 2 <= t) ++bi;
    while ((bi * (bi + 1)) / 2 > t) --bi;
    const int bj = t - (bi * (bi + 1)) / 2;

    const int row0 = bi * BM;
    const int col0 = bj * BN;

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

    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;

    // Global load assignment: one float4 per thread per tile of A and B.
    const int aRow = tid >> 1;        // 0..127
    const int aCol = (tid & 1) * 4;   // 0 or 4
    const int bRow = tid >> 5;        // 0..7
    const int bCol = (tid & 31) * 4;  // 0..124

    const float* Ag = A + (long)row0 * N + (long)aRow * N + aCol;
    const float* Bg = B + (long)bRow * N + col0 + bCol;

    // A[i][k] = 0 for k > i and B[k][j] = 0 for k < j, so the k-tiles of this
    // C-tile span only [bj*BM, (bi+1)*BM).
    const int ktStart = bj * (BM / BK);
    const int ktEnd   = (bi + 1) * (BM / BK);

    for (int kt = ktStart; kt < ktEnd; ++kt) {
        const int k0 = kt * BK;

        // Load A tile (BM x BK), store transposed into shared memory.
        float4 av = *reinterpret_cast<const float4*>(Ag + k0);
        As[aCol + 0][aRow] = av.x;
        As[aCol + 1][aRow] = av.y;
        As[aCol + 2][aRow] = av.z;
        As[aCol + 3][aRow] = av.w;

        // Load B tile (BK x BN) into shared memory.
        float4 bv = *reinterpret_cast<const float4*>(Bg + (long)k0 * N);
        *reinterpret_cast<float4*>(&Bs[bRow][bCol]) = bv;

        __syncthreads();

#pragma unroll
        for (int kk = 0; kk < BK; ++kk) {
            float4 a0 = *reinterpret_cast<const float4*>(&As[kk][threadRow * TM]);
            float4 a1 = *reinterpret_cast<const float4*>(&As[kk][threadRow * TM + 4]);
            float4 b0 = *reinterpret_cast<const float4*>(&Bs[kk][threadCol * TN]);
            float4 b1 = *reinterpret_cast<const float4*>(&Bs[kk][threadCol * TN + 4]);
            const float af[TM] = {a0.x, a0.y, a0.z, a0.w, a1.x, a1.y, a1.z, a1.w};
            const float bf[TN] = {b0.x, b0.y, b0.z, b0.w, b1.x, b1.y, b1.z, b1.w};
#pragma unroll
            for (int i = 0; i < TM; ++i)
#pragma unroll
                for (int j = 0; j < TN; ++j)
                    acc[i][j] = fmaf(af[i], bf[j], acc[i][j]);
        }
        __syncthreads();
    }

    // Fused tril: write zeros above the main diagonal (only diagonal tiles
    // actually need the mask; strictly lower tiles are fully valid).
    const bool full = (bi > bj);
    const int cBase = col0 + threadCol * TN;
#pragma unroll
    for (int i = 0; i < TM; ++i) {
        const int r = row0 + threadRow * TM + i;
        float* out = C + (long)r * N + cBase;
        if (full) {
            *reinterpret_cast<float4*>(out)     = make_float4(acc[i][0], acc[i][1], acc[i][2], acc[i][3]);
            *reinterpret_cast<float4*>(out + 4) = make_float4(acc[i][4], acc[i][5], acc[i][6], acc[i][7]);
        } else {
            float v[TN];
#pragma unroll
            for (int j = 0; j < TN; ++j)
                v[j] = (r >= cBase + j) ? acc[i][j] : 0.0f;
            *reinterpret_cast<float4*>(out)     = make_float4(v[0], v[1], v[2], v[3]);
            *reinterpret_cast<float4*>(out + 4) = make_float4(v[4], v[5], v[6], v[7]);
        }
    }
}

torch::Tensor lower_tri_mm_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() == at::kFloat && B.scalar_type() == at::kFloat,
                "inputs must be float32");
    TORCH_CHECK(A.dim() == 2 && A.size(0) == A.size(1), "A must be square");
    TORCH_CHECK(B.dim() == 2 && B.size(0) == B.size(1), "B must be square");
    TORCH_CHECK(A.size(0) == B.size(0), "A and B must have the same size");
    const int N = (int)A.size(0);
    TORCH_CHECK(N % BM == 0, "N must be divisible by 128");

    auto Ac = A.contiguous();
    auto Bc = B.contiguous();
    // zeros_like also clears the (never launched) upper-triangular tiles of C.
    auto C = torch::zeros_like(Ac);

    const int ntp = N / BM;
    const int blocks = ntp * (ntp + 1) / 2;
    lower_tri_mm_kernel<<<blocks, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
        Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(), N);
    return C;
}
"""

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

_lower_tri_mm = load_inline(
    name="lower_tri_mm",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["lower_tri_mm_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: replaces torch.matmul + torch.tril with a single fused
    custom CUDA kernel that exploits the lower-triangular structure to skip
    ~5.5x of the multiply-adds and applies the tril mask on write.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._op = _lower_tri_mm

    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.size(0) == A.size(1) and B.size(0) == B.size(1)
                and A.size(0) == B.size(0) and A.size(0) % 128 == 0):
            return self._op.lower_tri_mm_cuda(A, B)
        return torch.tril(torch.matmul(A, B))
```