```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 extension implementing blocked lower-triangular matmul.
#
# For lower-triangular A, B the product C = tril(A @ B) has the block sum
#     C[i, j] = sum_{k=j}^{i} A[i, k] @ B[k, j]
# so only block-column panels need to be computed:
#     C[k:, k:k2]  = A[k:, k:k2] @ B[k:k2, k:k2]     (beta = 0, first touch)
#     C[k:, :k]   += A[k:, k:k2] @ B[k:k2, :k]       (beta = 1, accumulate)
# This performs ~1.44 * N^3/3 FLOPs instead of 2 * N^3 (~4.2x fewer) while
# still using cuBLAS (TF32 tensor cores) for every block. Because the
# first-touch GEMMs use beta = 0, the output can be allocated with
# torch::empty (no full N^2 memset), and a small custom CUDA kernel zeroes
# only the strict upper triangle afterwards.
_CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <algorithm>

__global__ void zero_upper_kernel(float* __restrict__ c, const long long n) {
    const long long col = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (col >= n) return;
    for (long long row = blockIdx.y; row < col; row += gridDim.y) {
        c[row * n + col] = 0.0f;
    }
}

torch::Tensor trimul_lower(torch::Tensor A, torch::Tensor B, int64_t block_size) {
    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,
                "only float32 is supported");
    TORCH_CHECK(A.dim() == 2 && A.size(0) == A.size(1), "A must be square");
    TORCH_CHECK(B.dim() == 2 && B.size(0) == A.size(0) && B.size(1) == A.size(1),
                "B must have the same shape as A");
    TORCH_CHECK(block_size > 0, "block_size must be positive");

    const c10::cuda::CUDAGuard device_guard(A.device());
    auto Ac = A.contiguous();
    auto Bc = B.contiguous();
    const int64_t n = Ac.size(0);
    auto C = at::empty({n, n}, Ac.options());
    if (n == 0) {
        return C;
    }
    const int64_t bs = std::min<int64_t>(block_size, n);

    for (int64_t k = 0; k < n; k += bs) {
        const int64_t k2 = std::min(k + bs, n);
        auto Av = Ac.slice(0, k, n).slice(1, k, k2);        // (n-k) x (k2-k)
        auto Cv = C.slice(0, k, n).slice(1, k, k2);         // (n-k) x (k2-k)
        auto Bv = Bc.slice(0, k, k2).slice(1, k, k2);       // (k2-k) x (k2-k)
        Cv.addmm_(Av, Bv, /*beta=*/0, /*alpha=*/1);         // first touch
        if (k > 0) {
            auto Cv_l = C.slice(0, k, n).slice(1, 0, k);    // (n-k) x k
            auto Bv_l = Bc.slice(0, k, k2).slice(1, 0, k);  // (k2-k) x k
            Cv_l.addmm_(Av, Bv_l, /*beta=*/1, /*alpha=*/1); // accumulate
        }
    }

    if (n > 1) {
        cudaStream_t stream = at::cuda::getCurrentCUDAStream();
        const dim3 block(256);
        const dim3 grid((unsigned)((n + (int64_t)block.x - 1) / (int64_t)block.x),
                        (unsigned)std::min<int64_t>(n, 65535));
        zero_upper_kernel<<<grid, block, 0, stream>>>(C.data_ptr<float>(), (long long)n);
    }
    return C;
}
"""

_CPP_SRC = "torch::Tensor trimul_lower(torch::Tensor A, torch::Tensor B, int64_t block_size);"

try:
    _ext = load_inline(
        name="trimul_lower_ext",
        cpp_sources=_CPP_SRC,
        cuda_sources=_CUDA_SRC,
        functions=["trimul_lower"],
        verbose=False,
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3"],
    )
except Exception:
    _ext = None


class ModelNew(nn.Module):
    """
    Optimized C = tril(A @ B) for lower-triangular A, B.

    Uses a blocked decomposition evaluated with cuBLAS (TF32 tensor cores,
    as allowed) plus a custom CUDA kernel that zeroes the strict upper
    triangle. Falls back to the reference computation on unsupported inputs.
    """

    def __init__(self, block_size: int = 512) -> None:
        super().__init__()
        self.block_size = int(block_size)

    def forward(self, A, B):
        if (
            _ext is not None
            and A.is_cuda and B.is_cuda
            and A.dtype == torch.float32 and B.dtype == torch.float32
            and A.dim() == 2 and A.size(0) == A.size(1)
            and B.dim() == 2 and B.size(0) == A.size(0) and B.size(1) == A.size(1)
        ):
            prev_tf32 = torch.backends.cuda.matmul.allow_tf32
            torch.backends.cuda.matmul.allow_tf32 = True
            try:
                return _ext.trimul_lower(A, B, self.block_size)
            finally:
                torch.backends.cuda.matmul.allow_tf32 = prev_tf32
        return torch.tril(torch.matmul(A, B))
```