```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 matmul C = tril(A @ B).
#
# Exploits the triangular structure for large savings:
#  1. Only output blocks on/below the diagonal are computed (upper blocks are zero).
#  2. For each output block, the k reduction is restricted to [bx*BM, (by+1)*BM)
#     because A[r][k]=0 for k>r and B[k][c]=0 for k<c. This cuts total FLOPs ~6x
#     versus a dense matmul, and the separate tril kernel is eliminated entirely.
#
# The kernel uses 128x128 register-blocked tiles (8x8 per thread), vectorized
# float4 global/shared memory access, and padded shared memory to avoid bank conflicts.
tril_matmul_cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

#define BM 128
#define BN 128
#define BK 16
#define LDS 132  // 128 + 4 floats padding: keeps float4 alignment, avoids bank conflicts

__global__ __launch_bounds__(256)
void tril_matmul_kernel(const float* __restrict__ A,
                        const float* __restrict__ B,
                        float* __restrict__ C,
                        int N) {
    const int bx = blockIdx.x;  // column tile
    const int by = blockIdx.y;  // row tile
    if (bx > by) return;        // above the diagonal: C is zero (pre-initialized)

    __shared__ float As[BK][LDS];  // As[k][m] = A[by*BM + m][kt + k]
    __shared__ float Bs[BK][LDS];  // Bs[k][n] = B[kt + k][bx*BN + n]

    const int tx  = threadIdx.x;   // 0..15
    const int ty  = threadIdx.y;   // 0..15
    const int tid = ty * 16 + tx;  // 0..255

    // Each thread accumulates an 8x8 sub-tile, arranged as 2x2 groups of 4x4
    // so that shared memory float4 accesses are bank-conflict free.
    float acc[2][2][4][4] = {};

    const int k_begin = bx * BM;
    const int k_end   = min((by + 1) * BM, N);
    const bool aligned = ((N & 3) == 0);

    // Thread -> load assignments
    const int a_row = tid >> 2;   // A: 128 rows x 16 cols, 4 float4 per row
    const int a_kv  = tid & 3;
    const int b_k   = tid >> 5;   // B: 16 rows x 128 cols, 32 float4 per row
    const int b_nv  = tid & 31;

    for (int kt = k_begin; kt < k_end; kt += BK) {
        // ---- Load A tile into shared (transposed layout) ----
        #pragma unroll
        for (int p = 0; p < 2; ++p) {
            const int m  = a_row + p * 64;
            const int gm = by * BM + m;
            const int gc = kt + a_kv * 4;
            float v0 = 0.f, v1 = 0.f, v2 = 0.f, v3 = 0.f;
            if (gm < N) {
                if (aligned && gc + 3 < N) {
                    const float4 v = *reinterpret_cast<const float4*>(A + (long)gm * N + gc);
                    v0 = v.x; v1 = v.y; v2 = v.z; v3 = v.w;
                } else {
                    const float* base = A + (long)gm * N;
                    v0 = (gc + 0 < N) ? base[gc + 0] : 0.f;
                    v1 = (gc + 1 < N) ? base[gc + 1] : 0.f;
                    v2 = (gc + 2 < N) ? base[gc + 2] : 0.f;
                    v3 = (gc + 3 < N) ? base[gc + 3] : 0.f;
                }
            }
            As[a_kv * 4 + 0][m] = v0;
            As[a_kv * 4 + 1][m] = v1;
            As[a_kv * 4 + 2][m] = v2;
            As[a_kv * 4 + 3][m] = v3;
        }
        // ---- Load B tile into shared ----
        #pragma unroll
        for (int p = 0; p < 2; ++p) {
            const int k  = b_k + p * 8;
            const int gk = kt + k;
            const int gn = bx * BN + b_nv * 4;
            float v0 = 0.f, v1 = 0.f, v2 = 0.f, v3 = 0.f;
            if (gk < N) {
                if (aligned && gn + 3 < N) {
                    const float4 v = *reinterpret_cast<const float4*>(B + (long)gk * N + gn);
                    v0 = v.x; v1 = v.y; v2 = v.z; v3 = v.w;
                } else {
                    const float* base = B + (long)gk * N;
                    v0 = (gn + 0 < N) ? base[gn + 0] : 0.f;
                    v1 = (gn + 1 < N) ? base[gn + 1] : 0.f;
                    v2 = (gn + 2 < N) ? base[gn + 2] : 0.f;
                    v3 = (gn + 3 < N) ? base[gn + 3] : 0.f;
                }
            }
            *reinterpret_cast<float4*>(&Bs[k][b_nv * 4]) = make_float4(v0, v1, v2, v3);
        }
        __syncthreads();

        // ---- Compute 8x8 per thread ----
        #pragma unroll
        for (int k = 0; k < BK; ++k) {
            float a_frag[2][4];
            float b_frag[2][4];
            #pragma unroll
            for (int g = 0; g < 2; ++g) {
                const float4 v = *reinterpret_cast<const float4*>(&As[k][g * 64 + ty * 4]);
                a_frag[g][0] = v.x; a_frag[g][1] = v.y; a_frag[g][2] = v.z; a_frag[g][3] = v.w;
            }
            #pragma unroll
            for (int g = 0; g < 2; ++g) {
                const float4 v = *reinterpret_cast<const float4*>(&Bs[k][g * 64 + tx * 4]);
                b_frag[g][0] = v.x; b_frag[g][1] = v.y; b_frag[g][2] = v.z; b_frag[g][3] = v.w;
            }
            #pragma unroll
            for (int gi = 0; gi < 2; ++gi)
                #pragma unroll
                for (int gj = 0; gj < 2; ++gj)
                    #pragma unroll
                    for (int i = 0; i < 4; ++i)
                        #pragma unroll
                        for (int j = 0; j < 4; ++j)
                            acc[gi][gj][i][j] += a_frag[gi][i] * b_frag[gj][j];
        }
        __syncthreads();
    }

    // ---- Store results (vectorized for fully-interior blocks below the diagonal) ----
    const bool fast = (bx < by) && ((by + 1) * BM <= N) && aligned;
    #pragma unroll
    for (int gi = 0; gi < 2; ++gi) {
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            const int r = by * BM + gi * 64 + ty * 4 + i;
            #pragma unroll
            for (int gj = 0; gj < 2; ++gj) {
                const int c0 = bx * BN + gj * 64 + tx * 4;
                if (fast) {
                    *reinterpret_cast<float4*>(C + (long)r * N + c0) =
                        make_float4(acc[gi][gj][i][0], acc[gi][gj][i][1],
                                    acc[gi][gj][i][2], acc[gi][gj][i][3]);
                } else if (r < N) {
                    #pragma unroll
                    for (int j = 0; j < 4; ++j) {
                        const int c = c0 + j;
                        if (c < N && c <= r) C[(long)r * N + c] = acc[gi][gj][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,
                "inputs must be float32");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2 && 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({N, N}, Ac.options());
    if (N == 0) return C;

    const int nb = (N + BM - 1) / BM;
    dim3 grid(nb, nb);
    dim3 block(16, 16);
    tril_matmul_kernel<<<grid, block>>>(
        Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(), N);
    return C;
}
"""

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

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


class ModelNew(nn.Module):
    """
    Optimized model: single fused CUDA kernel computing C = tril(A @ B) for
    lower triangular A and B. Skips all-zero upper-triangular blocks and
    restricts the reduction range, eliminating ~5/6 of the dense-matmul FLOPs
    and the separate tril 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)
```