```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

# Blocked lower-triangular matmul:
#
#   For lower-triangular A, B the product C = tril(A @ B) satisfies
#       C[i, j] = sum_{s=j}^{i} A[i, s] * B[s, j],
#   so, with column blocks of width bs, only the following block-GEMMs are
#   required (every strictly above-diagonal block is exactly zero):
#
#       C[k:, :k2] += A[k:, k:k2] @ B[k:k2, :k2]        k = 0, bs, 2*bs, ...
#
#   This performs ~4x fewer FLOPs than the dense matmul + tril reference
#   while still running entirely on cuBLAS TF32 tensor-core GEMMs (TF32 is
#   allowed, as in the PyTorch reference).  Details that make it fast:
#     * the loop runs in a single C++ call with raw cublasGemmEx calls
#       (no per-iteration Python slicing / dispatcher overhead);
#     * block k = 0 uses beta = 0 (first touch of its panel), so only
#       columns [bs:, N) need pre-zeroing;
#     * that cudaMemset2DAsync runs on a side stream, overlapped with the
#       first GEMM (they touch disjoint columns of C);
#     * no tril / zero-upper pass is needed afterwards: off-diagonal
#       *blocks* are never written (they are zeroed by the memset), and
#       the strict upper triangle inside each diagonal block is written
#       by its GEMM with exact 0.0f values (every product term has a zero
#       factor).
_CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cublas_v2.h>
#include <cuda_runtime.h>
#include <algorithm>

#define CUBLAS_CHECK(expr)                                          \
  do {                                                              \
    cublasStatus_t _st = (expr);                                    \
    TORCH_CHECK(_st == CUBLAS_STATUS_SUCCESS,                       \
                "cuBLAS error ", (int)_st, " at ", __FILE__, ":",   \
                __LINE__);                                          \
  } while (0)

#define CUDA_CHECK(expr)                                            \
  do {                                                              \
    cudaError_t _err = (expr);                                      \
    TORCH_CHECK(_err == cudaSuccess, "CUDA error: ",                \
                cudaGetErrorString(_err));                          \
  } while (0)

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");

    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::max<int64_t>(1, std::min(block_size, n));

    float* __restrict__ Cp = C.data_ptr<float>();
    const float* __restrict__ Ap = Ac.data_ptr<float>();
    const float* __restrict__ Bp = Bc.data_ptr<float>();

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    CUBLAS_CHECK(cublasSetStream(handle, stream));

    // TF32 tensor-core math on sm80+, plain FP32 otherwise.
    auto* prop = at::cuda::getCurrentDeviceProperties();
    const cublasComputeType_t ctype =
        (prop->major >= 8) ? CUBLAS_COMPUTE_32F_FAST_TF32 : CUBLAS_COMPUTE_32F;

    const size_t elt = sizeof(float);

    // Pre-zero C[:, bs:] (columns written with beta=1 later).  This runs on a
    // side stream so it overlaps with the first (beta=0) GEMM, which only
    // touches C[:, :bs].  Block 0's panel is fully overwritten, so it never
    // needs zero-initialization.
    cudaStream_t zero_stream = nullptr;
    cudaEvent_t fork_ev, join_ev;
    bool have_side = (bs < n);
    if (have_side) {
        CUDA_CHECK(cudaEventCreateWithFlags(&fork_ev, cudaEventDisableTiming));
        CUDA_CHECK(cudaEventCreateWithFlags(&join_ev, cudaEventDisableTiming));
        zero_stream = at::cuda::getStreamFromPool(/*isHighPriority=*/false,
                                                  A.device().index());
        CUDA_CHECK(cudaEventRecord(fork_ev, stream));
        CUDA_CHECK(cudaStreamWaitEvent(zero_stream, fork_ev, 0));
        CUDA_CHECK(cudaMemset2DAsync(Cp + bs, n * elt, 0, (n - bs) * elt, n,
                                     zero_stream));
        CUDA_CHECK(cudaEventRecord(join_ev, zero_stream));
    }

    const float alpha = 1.0f;
    for (int64_t k = 0; k < n; k += bs) {
        const int64_t k2 = std::min(k + bs, n);
        if (k == bs && have_side) {
            // Ensure the pre-zeroing finished before the first beta=1 GEMM.
            CUDA_CHECK(cudaStreamWaitEvent(stream, join_ev, 0));
        }
        const float beta = (k == 0) ? 0.0f : 1.0f;
        const float* Asub = Ap + k * n + k;   // (n-k) x (k2-k), ld = n
        const float* Bsub = Bp + k * n;       // (k2-k) x k2,    ld = n
        float* Csub = Cp + k * n;             // (n-k) x k2,     ld = n
        // Row-major C = A @ B  <=>  column-major C' = B' @ A'.
        CUBLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N,
                                  (int)k2, (int)(n - k), (int)(k2 - k),
                                  &alpha,
                                  Bsub, CUDA_R_32F, (int)n,
                                  Asub, CUDA_R_32F, (int)n,
                                  &beta,
                                  Csub, CUDA_R_32F, (int)n,
                                  ctype, CUBLAS_GEMM_DEFAULT));
    }

    if (have_side) {
        CUDA_CHECK(cudaEventRecord(fork_ev, stream));
        CUDA_CHECK(cudaStreamWaitEvent(zero_stream, fork_ev, 0));
        CUDA_CHECK(cudaEventDestroy(join_ev));
        CUDA_CHECK(cudaEventDestroy(fork_ev));
    }
    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_cublas_ext",
        cpp_sources=_CPP_SRC,
        cuda_sources=_CUDA_SRC,
        functions=["trimul_lower"],
        verbose=False,
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3"],
        extra_ldflags=["-lcublas", "-lcudart"],
    )
except Exception:
    _ext = None


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

    Evaluates the product as ~N/(2*bs) cuBLAS (TF32 tensor-core) block-GEMMs
    covering only the lower block triangle (~4x fewer FLOPs than dense
    matmul + tril), with the output pre-zeroing overlapped on a side CUDA
    stream.  Falls back to the reference computation for 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)
        ):
            return _ext.trimul_lower(A, B, self.block_size)
        return torch.tril(torch.matmul(A, B))
```