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

# Multi-stream 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 (panels) of width bs, panel j satisfies
#       C[k*bs:, panel j] += A[k*bs:, k-block k] @ B[k-block k, panel j],
#   for k = j, j+1, ..., and every strictly above-diagonal block is
#   exactly zero.  Each panel accumulates into its own columns only, so
#   different panels are INDEPENDENT and can run concurrently on separate
#   CUDA streams; the k-loop within a panel is sequential (beta = 0 for
#   k == j, then beta = 1).
#
#   Compared to the single-stream version this keeps the GPU saturated:
#   per-GEMM wave-quantization tails of one panel are filled with tiles of
#   the other panels, while total FLOPs stay ~4x below the dense matmul +
#   tril reference.  Details:
#     * one C++ call issues all ~nb(nb+1)/2 raw cublasGemmEx calls (TF32
#       tensor-core math, as allowed by the reference), with no Python or
#       dispatcher overhead in the loop;
#     * every panel stream gets its own cuBLAS workspace so concurrent
#       GEMMs never race on shared scratch memory;
#     * the strictly-above-diagonal blocks are zeroed by cudaMemset2DAsync
#       calls on an extra stream, overlapped with the compute; the strict
#       upper triangle inside each diagonal block is written by its first
#       (beta=0) GEMM with exact 0.0f values (every product term has a zero
#       factor), so no tril pass is needed afterwards;
#     * events fork all worker streams from the caller's stream and join
#       them back before returning.
_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>
#include <mutex>
#include <unordered_map>
#include <vector>

#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)

namespace {

constexpr size_t kWorkspaceSize = size_t(8) << 20;  // 8 MiB per stream
constexpr int kMaxWorkers = 16;

struct WorkerPool {
    std::vector<cudaStream_t> streams;
    std::vector<cudaEvent_t> events;   // per-stream "done" events
    std::vector<void*> workspaces;
    cudaEvent_t fork = nullptr;
    cudaEvent_t zero_done = nullptr;
    cudaStream_t zero_stream = nullptr;
};

std::unordered_map<int, WorkerPool> g_pools;
std::once_flag g_events_flag;

void init_shared_events() {
    // Create the reusable fork / join events on the current device context.
    // They are re-created per device lazily inside get_pool(); this flag only
    // keeps the code path simple (events are cheap and created once per pool).
}

WorkerPool& get_pool(int device, int needed) {
    auto it = g_pools.find(device);
    if (it == g_pools.end()) {
        WorkerPool pool;
        const c10::cuda::CUDAGuard guard(device);
        CUDA_CHECK(cudaEventCreateWithFlags(&pool.fork, cudaEventDisableTiming));
        CUDA_CHECK(cudaEventCreateWithFlags(&pool.zero_done, cudaEventDisableTiming));
        it = g_pools.emplace(device, std::move(pool)).first;
    }
    WorkerPool& pool = it->second;
    const c10::cuda::CUDAGuard guard(device);
    while ((int)pool.streams.size() < needed) {
        cudaStream_t s = at::cuda::getStreamFromPool(/*isHighPriority=*/false,
                                                     device).stream();
        cudaEvent_t e;
        CUDA_CHECK(cudaEventCreateWithFlags(&e, cudaEventDisableTiming));
        void* ws = nullptr;
        CUDA_CHECK(cudaMalloc(&ws, kWorkspaceSize));
        pool.streams.push_back(s);
        pool.events.push_back(e);
        pool.workspaces.push_back(ws);
    }
    if (pool.zero_stream == nullptr) {
        pool.zero_stream = at::cuda::getStreamFromPool(/*isHighPriority=*/false,
                                                       device).stream();
    }
    return pool;
}

}  // namespace

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>(16, std::min(block_size, n));
    const int nb = (int)((n + bs - 1) / bs);

    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 main_stream = at::cuda::getCurrentCUDAStream();
    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();

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

    const int device = A.device().index();
    const int S = std::min(std::max(nb, 1), kMaxWorkers);
    WorkerPool& pool = get_pool(device, S);

    // Fork: every worker stream waits until A, B and C are ready.
    CUDA_CHECK(cudaEventRecord(pool.fork, main_stream));

    // Zero the strictly-above-diagonal blocks of C on a side stream,
    // overlapped with the panel GEMMs (they never write there).
    cudaStream_t zero_stream = pool.zero_stream;
    const bool have_zero = (nb > 1);
    if (have_zero) {
        CUDA_CHECK(cudaStreamWaitEvent(zero_stream, pool.fork, 0));
        const size_t elt = sizeof(float);
        for (int j = 1; j < nb; ++j) {
            const int64_t c0 = (int64_t)j * bs;            // first col of panel j
            const int64_t width = std::min(bs, n - c0);    // panel width
            CUDA_CHECK(cudaMemset2DAsync(Cp + c0, n * elt, 0, width * elt, c0,
                                         zero_stream));
        }
        CUDA_CHECK(cudaEventRecord(pool.zero_done, zero_stream));
    }

    const float one = 1.0f;
    const float zero = 0.0f;
    std::vector<cudaEvent_t> used_events;
    used_events.reserve(S);

    // Launch panels, each on its own stream.  Worker s takes panels
    // j = s, s+S, s+2S, ... (only relevant when nb > kMaxWorkers).
    for (int j0 = 0; j0 < std::min(nb, S); ++j0) {
        cudaStream_t s = pool.streams[j0];
        void* ws = pool.workspaces[j0];
        CUDA_CHECK(cudaStreamWaitEvent(s, pool.fork, 0));
        CUBLAS_CHECK(cublasSetStream(handle, s));
        CUBLAS_CHECK(cublasSetWorkspace(handle, ws, kWorkspaceSize));
        for (int j = j0; j < nb; j += S) {
            const int64_t pc = (int64_t)j * bs;            // panel first col
            const int64_t pw = std::min(bs, n - pc);       // panel width
            const float beta0 = 0.0f;
            for (int k = j; k < nb; ++k) {
                const int64_t kr = (int64_t)k * bs;        // k-block first row
                const int64_t kd = std::min(bs, n - kr);   // k-block depth
                const int64_t rows = n - kr;               // rows of C updated
                const float* Asub = Ap + kr * n + kr;      // rows x kd, ld = n
                const float* Bsub = Bp + kr * n + pc;      // kd x pw,  ld = n
                float* Csub = Cp + kr * n + pc;            // rows x pw, ld = n
                // Row-major Csub = Asub @ Bsub  <=>
                // column-major Csub' = Bsub' @ Asub'.
                CUBLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N,
                                          (int)pw, (int)rows, (int)kd,
                                          &one,
                                          Bsub, CUDA_R_32F, (int)n,
                                          Asub, CUDA_R_32F, (int)n,
                                          (k == j) ? &beta0 : &one,
                                          Csub, CUDA_R_32F, (int)n,
                                          ctype, CUBLAS_GEMM_DEFAULT));
            }
        }
        CUDA_CHECK(cudaEventRecord(pool.events[j0], s));
        used_events.push_back(pool.events[j0]);
    }

    // Restore the caller's stream on the cuBLAS handle (the main stream keeps
    // whatever workspace PyTorch assigned to it).
    CUBLAS_CHECK(cublasSetStream(handle, main_stream));

    // Join: the returned tensor is only ready once every worker and the
    // zeroing stream are done.
    for (cudaEvent_t e : used_events) {
        CUDA_CHECK(cudaStreamWaitEvent(main_stream, e, 0));
    }
    if (have_zero) {
        CUDA_CHECK(cudaStreamWaitEvent(main_stream, pool.zero_done, 0));
    }
    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_ms_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 cuBLAS TF32 tensor-core block-GEMMs covering
    only the lower block triangle (~4x fewer FLOPs than dense matmul +
    tril), with independent column panels executed concurrently on
    multiple CUDA streams (each with its own cuBLAS workspace) and the
    output pre-zeroing overlapped on a side 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))
```